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

We thank the reviewer for their thoughtful comments. Below, we address the identified weaknesses and questions.

"Although this is a theoretical paper, it lacks any form of empirical validation of its conclusions, even in a simplified format."

While we did not conduct numerical simulations, numerous experiments in the literature demonstrate the effectiveness of U-Nets in image-denoising tasks and hierarchical denoising models, including [KGMS23] and [SFW24]. We believe these existing experiments provide good empirical validation for our theoretical framework.

We acknowledge the importance of experimental validation to make the paper more self-contained and regret that we may not be able to complete simulations within the short term. Nonetheless, we plan to include experimental results in the final version.

[KGMS23] Zahra Kadkhodaie, Florentin Guth, Stephane Mallat, and Eero P Simoncelli. Learning multi-scale local conditional probability models of images. ICLR 2023.

[SFW24] Antonio Sclocchi, Alessandro Favero, and Matthieu Wyart. A phase transition in diffusion models reveals the hierarchical nature of data. arXiv preprint arXiv:2402.16991, 2024.

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

We agree that ResNets (with short skip connections, assuming all intermediate layer dimensions are kept the same) and transformers could also simulate the belief propagation process and serve as effective denoising functions. In principle, we could theoretically demonstrate that ResNets and transformers can approximate belief propagation in the GHM framework, though they might require more parameters than U-Nets. That said, U-Nets are not the only architectures capable of performing well on denoising tasks.

However, ResNets and transformers may involve redundant operations that could make them less efficient than U-Nets in hierarchical models, given the same number of training samples. Intuitively, U-Nets represent a minimal structure that effectively parameterizes the belief propagation process, making them particularly well-suited for such tasks.

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

We apologize for the confusion caused by the title. Our intention was not to suggest that U-Nets are efficient for classification. Instead, we aimed to demonstrate that ConvNets are efficient for classification tasks, while U-Nets are efficient for denoising and diffusion tasks.

I appreciate the authors' feedback. It is understandable that the authors were unable to include toy experiments to quantify the impact of the sample complexity advantages at this stage. However, as promised, please ensure these results are included in the final version. Additionally, kindly update the text to clarify the connection between the UNet and classification to avoid confusion. Beyond these points, I believe the authors have sufficiently addressed my concerns, and I am inclined to raise the score to a 6.

We thank the reviewer for supporting our paper. Below, we address the identified weaknesses and questions.

"The main weakness of the paper is that the data assumptions are very restrictive. … 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). "

Belief propagation is indeed possible and efficient as long as the number of children per node is small. The conditional independence assumption serves as a simplification and proof of concept. We believe that even without this assumption, through a more careful analysis of approximation theory, the denoising function would still admit efficient U-Net approximation.

"In particular, they exist simpler and more efficient classical algorithms to learn and sample from these distributions than diffusion models. "

We agree that simpler and more efficient classical algorithms exist for learning and sampling from such distributions. For instance, a classical approach could involve using a maximum likelihood estimator to learn the graphical model and then running belief propagation on the learned model for denoising.

Our theoretical contribution demonstrates that diffusion models can also perform the same task efficiently. Conceptually, diffusion models have the advantage of relying less on the precise specification of the graphical model. Classical approaches may suffer when the model is misspecified, whereas diffusion models learn the score function directly using neural networks without explicit knowledge of the data-generating process. This versatility makes diffusion models particularly appealing for applications where the data-generating process is unknown.

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?

We believe the GHM captures some key properties of image distributions. There is existing work demonstrating the effectiveness of hierarchical models (not necessarily with the conditional independence assumption) in modeling image distributions, such as [KGMS23], [SFW24], and [PCTFW23].

However, the GHM proposed in this paper has limitations, which could be addressed in future work: (1) The conditional independence assumption is restrictive and could be generalized to better reflect real-world data. (2) Each hidden node in the current model is assumed to be a discrete variable taking values in [S]. Extending the framework to include continuous variables would be an interesting and valuable direction.

[KGMS23] Zahra Kadkhodaie, Florentin Guth, Stephane Mallat, and Eero P Simoncelli. Learning multi-scale local conditional probability models of images. ICLR 2023.

[SFW24] Antonio Sclocchi, Alessandro Favero, and Matthieu Wyart. A phase transition in diffusion models reveals the hierarchical nature of data. arXiv preprint arXiv:2402.16991, 2024.

[PCTFW23] Leonardo Petrini, Francesco Cagnetta, Umberto M Tomasini, Alessandro Favero, and Matthieu Wyart. How deep neural networks learn compositional data: The random hierarchy model. arXiv preprint arXiv:2307.02129, 2023.

We thank the reviewer for their thoughtful comments. Below, we address the identified weaknesses and questions.

"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. Can the authors clarify whether true convolution or cross-correlation is used, given that parameter sharing and local invariance properties are retained? "

The convolution operation used in real ConvNets differs from the one in this paper in the following way: in standard ConvNets, each image patch undergoes an inner product with a convolution filter (a linear transformation of each pixel and aggregation of pixel information), followed by a ReLU activation function (pointwise nonlinearity). In contrast, in this paper, each image pixel is first processed by a nonlinear transformation (multiple linear transformations and pointwise nonlinearities) before the pixel information is aggregated. The key distinction lies in the order of operations—“linear transformation”, “pointwise nonlinearity”, and “pixel information aggregation”, while translation invariance is preserved.

After submitting the paper, we realized that multiple real convolutional layers could simulate the convolutional layer used in our paper. Specifically, each real convolution filter would need to have only one vector of non-zero weights corresponding to a single pixel, with all other pixel weights set to zero. While this pattern is not typically observed in practice, it is theoretically valid and can simulate the pixel-wise linear transformation employed in our paper. (We also notice that ​​in real U-Nets, people indeed use multiple convolutional layers at each level) We will include this clarification in the final version of the paper.

"The practical utility of the provided bounds is unclear, as their direct practical applicability for empirical settings is not fully established. 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?"

The theoretical bound avoids the curse of dimensionality in neural network approximation, demonstrating the efficiency of neural networks in hierarchical models. While the bound itself may not have immediate practical applications, the intuition provided by the theoretical results could be insightful.

In proving the theoretical results, we leverage the fact that each layer of the U-Net approximates a belief-propagation step in the denoising process. This suggests a hypothesis that the intermediate layers of U-Nets, after training, represent intermediate-level concepts within a hierarchical model. This hypothesis is supported by several works, including [KGMS23] and [SFW23]. Future research could explore this observation for algorithm design. For instance, the theoretical construction suggests that the middle layers may represent high-level image concepts, such as “birds” or “dogs.” These intermediate representations could potentially serve as useful features for the image classification task.

[KGMS23] Zahra Kadkhodaie, Florentin Guth, Stephane Mallat, and Eero P Simoncelli. Learning multi-scale local conditional probability models of images. ICLR 2023.

[SFW24] Antonio Sclocchi, Alessandro Favero, and Matthieu Wyart. A phase transition in diffusion models reveals the hierarchical nature of data. arXiv preprint arXiv:2402.16991, 2024.

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

Regarding the alignment between theory and practice: as mentioned in response to the previous question, the theory predicts that the intermediate layers of U-Nets represent intermediate-level concepts of a transformer. This hypothesis is partially supported by several works that have conducted experimental studies of U-Nets, including [KGMS23] and [SFW23].

Some gaps between theory and practice remain. In most practical diffusion models, particularly those used for image generation, the models are conditional, with text input serving as guidance. Our current model does not incorporate the text input. Future work will explore joint modeling of both image and text distributions and investigate the architecture for joint belief propagation processes.

We thank the reviewer for their thoughtful comments. Below, we address the identified weaknesses and questions.

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

While we did not conduct numerical simulations, numerous experiments in the literature demonstrate the effectiveness of U-Nets in image-denoising tasks and hierarchical denoising models, including [KGMS23] and [SFW24]. We acknowledge the value of such experiments and regret that we may not be able to complete simulations in the short term. However, we plan to include experimental results in the final version.

[KGMS23] Zahra Kadkhodaie, Florentin Guth, Stephane Mallat, and Eero P Simoncelli. Learning multi-scale local conditional probability models of images. ICLR 2023.

[SFW24] Antonio Sclocchi, Alessandro Favero, and Matthieu Wyart. A phase transition in diffusion models reveals the hierarchical nature of data. arXiv preprint arXiv:2402.16991, 2024.

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

Eq. (17) analyzes ERM over the neural network class with a norm constraint, which is critical for proving the generalization error bound. In practice, minimizing the norm-constrained ERM can be equivalently reformulated as minimizing the norm-regularized ERM, such as via SGD on the regularized empirical risk. Therefore, adding a regularization term does not significantly affect the conclusions drawn from the analysis.

In neural network, the upper bound in the theorem has the tendency to be too loose, I was wondering how tight are the bounds?

We acknowledge that the upper bound in the theorem could be improved by adopting a more refined approximation scheme for the belief propagation algorithm. While the current bound may not be tight, it ensures that the required sample size scales polynomially with the dimension and relevant problem parameters, thereby avoiding the curse of dimensionality. This ability to circumvent the curse of dimensionality represents a key contribution of this paper.

We thank the reviewer for their thoughtful comments. Below, we address the identified weaknesses and questions.

``The goal of the analysis is slightly unclear… However, there is no motivation on why such a bound could be useful. Why are the established bounds useful for future neural network architecture development? Could they help us understand scaling laws for U-Nets?’’

The primary goal of this analysis is to address the challenge of the curse of dimensionality in neural network approximation theory: approximating general high-dimensional functions requires network sizes that scale exponentially with dimension $d$. Existing works mitigate the curse of dimensionality by leveraging smoothness or low-dimensional structure assumptions of the target function (e.g., [Bach17]). However, applying a smoothness assumption to the denoising function of image distributions imposes unnatural constraints on the density function in high-dimensional spaces, as noted in [OAS23].

This work introduces a different approach by assuming a generative hierarchical model (GHM). This model generates high-level concepts first and progressively refines them into finer details, aligning with recent advances in modeling image and text distributions (e.g., [SFW24]). By leveraging this assumption, the neural network naturally approximates the belief propagation algorithm within the GHM framework, circumventing the curse of dimensionality.

This framework offers a new statistical perspective on why U-Nets (and potentially other deep neural networks) perform effectively in denoising tasks and diffusion models. Understanding the broader implications of this theoretical framework is a key direction for future research, including its potential impact on architecture design and scaling laws.

[OAS23] Kazusato Oko, Shunta Akiyama, and Taiji Suzuki, Diffusion models are minimax optimal distribution estimators, arXiv preprint arXiv:2303.01861 (2023).

[Bach17] Francis Bach. Breaking the Curse of Dimensionality with Convex Neural Networks. JMLR 18(19):1−53, 2017.

[SFW24] Antonio Sclocchi, Alessandro Favero, and Matthieu Wyart. A phase transition in diffusion models reveals the hierarchical nature of data. arXiv preprint arXiv:2402.16991, 2024.

``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]? ‘’

Our approximation theory suggests that the intermediate layers of U-Nets correspond to the beliefs of latent nodes in the hierarchical model. Specifically: The top and bottom layers of U-Nets represent the most fine-grained features. The middle layers capture higher-level concepts, aligning with the hierarchical nature of the generative model.

We propose the hypothesis that after training U-Nets via stochastic gradient algorithms, the intermediate layers represent the beliefs of latent nodes in the generative hierarchical model. This hypothesis aligns with empirical observations in [SFW24], which demonstrate the hierarchical nature of data in diffusion models. This perspective may also explain why the semantic information in U-Net's intermediate representations is valuable for tasks like segmentation, as supported by results in [1]. We plan to further investigate this hypothesis and its implications for downstream tasks in future work.

[1] Baranchuk, Dmitry, et al. "Label-Efficient Semantic Segmentation with Diffusion Models." International Conference on Learning Representations.

[SFW24] Antonio Sclocchi, Alessandro Favero, and Matthieu Wyart. A phase transition in diffusion models reveals the hierarchical nature of data. arXiv preprint arXiv:2402.16991, 2024.