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

Authors: Song Mei
Primary area: learning theory
Keywords: diffusion model, denoising, belief propagation, generative hierarchical models

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

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

Decision: accept (poster)
Reviewer scores: 6, 6, 8, 6, 6

Reviews and discussion:
<thread>
creation date:
1730679138854

thread_type:
official_review

<thread_object>
creator: reviewer_5xwe

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

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.

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.

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730680382283

thread_type:
official_review

<thread_object>
creator: reviewer_xjiq

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

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.

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.

questions:
See the Weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730694452048

thread_type:
official_review

<thread_object>
creator: reviewer_ukfv

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

strengths:
This paper studies an important problem, which is to understand what properties of the data distribution coupled with inductive biases of network architectures allow to break the curse of dimensionality when learning diffusion models. The observation that CNNs and UNet correspond to classical Bayesian inference algorithms in the case of tree-structured models is an interesting one. To the best of my knowledge, this is the first work that obtains such results for deep networks.

weaknesses:
The main weakness of the paper is that the data assumptions are very restrictive. In particular, they exist simpler and more efficient classical algorithms to learn and sample from these distributions than diffusion models.

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730799307218

thread_type:
official_review

<thread_object>
creator: reviewer_f34u

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

strengths:
- Provide a rigorous theoretical explanation of U-Net and CNN through the lens of a graphical model

weaknesses:
- Many hypotheses are made for the derivation. The conclusion would be strengthened if there were at least some toy-ish experiments to test these hypotheses
- (minor) Given how dense the paper is, a lookup table for notations would be very helpful

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730913145065

thread_type:
official_review

<thread_object>
creator: reviewer_iuq6

<content>
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

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.

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.

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.
</content>
</thread_object>

</thread>


Metareview:
metareview:
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.