Weaknesses
- **inductive bias**: theorem 2 is the cornerstone of the theoretical explanation of this work as to why U-Nets are effective in diffusion. Yet I think that it has 2 problems. First, the result is about the Haar decomposition of the diffusion process, and therefore it applies to all models that have a multiscale structure with average pooling, not specifically U-Nets. Second, and in my view much more problematic, diffusion pipelines often use an "epsilon-type" prediction framework, in which the role of the U-Net is not to predict the denoised image (or signal in general), but rather the noise.
There are multiple examples of this that can be found either in the examples of the diffusers library of HuggingFace or in the seminal papers of Ho and Song:
- in the unconditional generation example: https://github.com/huggingface/diffusers/blob/main/examples/unconditional_image_generation/train_unconditional.py#L576, when prediction type is "epsilon", the loss is computed as the MSE between the output of the model and the noise, and one can see that this is the default mode (L229), not changed in the way the script is called (https://github.com/huggingface/diffusers/blob/main/examples/unconditional_image_generation/README.md)
- in the text-to-image generation example, the prediction type is taken by default from the noise scheduler which has it by default has "epsilon" (https://github.com/huggingface/diffusers/blob/716286f19ddd9eb417113e064b538706884c8e73/src/diffusers/schedulers/scheduling_ddpm.py#L121), not changed in the way the script is called (https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/README.md)
- In [7], algorithm 1 clearly shows that it's the noise that is learned
- In the implementation of [34], the noise `z` is used used in the MSE with the model output (https://github.com/yang-song/score_sde/blob/main/losses.py#L115)
As a consequence I think the theoretical result does not cover the practice. Indeed, it is said "This inductive bias enables the encoder and decoder networks to focus on the signal on a low enough frequency which is not dominated by noise." while the network's role is precisely to focus on subspaces which feature the noise in order to return it. While I understand that some pipelines still use a "sample" type prediction, I think that the understanding of the U-Net inductive bias is not understood if it doesn't take into account both types of predictions, especially if the explanation of one type is contrarian to the other type.
- **clarity and space use**: I think that some sections of the paper are very unclear, in part because some essential parts were left out in the appendix. For example, Algorithm 1 for staged training should be in the core paper. Similarly the practical construction of a U-Net that guarantees boundary conditions should be in the core paper. For now, we just get an idea of what the decoder subspaces should look like, but not how the decoder should be designed in order to achieve these subspaces. Right now it is not clear to me how these decoders are designed. A last example (although there are some more) is that the formal definition of the U-Net is hard to grasp at first read. There should be a correspondence between the elements introduced and the typical 2D image segmentation U-Net architecture.
- **prior work comparison/mention**: I noticed two parts that needed more references to prior work/comparison to it. First, the connection between U-Nets and wavelets has been a topic of many works among which [A, B]. Second, diffusion on complicated geometries including manifolds has been recently covered by works such as [C-F]. For both of these, mentioning as well as highlighting the differences would be super important.
*Minor*:
- To me preconditioning means finding a matrix close to the Hessian inverse in an optimization procedure in order to better take into account the geometry of the problem when optimizing, not finding a good initialization. Here it seems that preconditioning is more used to describe a good initialization rather than taking into account geometric information.
- "It also explains why in practice, encoders are often chosen significantly smaller than the decoder [26]". I am not sure about this fact, so I think it needs much more evidence (see for example in the diffusers library https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unet_2d.py#L198
the decoder only has one more layer than the encoder). Also [26] is about VAEs, not U-Nets.
[A]: Ramzi, Z., Michalewicz, K., Starck, JL. et al. Wavelets in the Deep Learning Era. J Math Imaging Vis 65, 240–251 (2023).
[B]: Liu, P., Zhang, H., Zhang, K., Lin, L., & Zuo, W. (2018). Multi-level wavelet-CNN for image restoration. In Proceedings of the IEEE conference on computer vision and pattern recognition workshops (pp. 773-782).
[C]: De Bortoli, V., Mathieu, E., Hutchinson, M., Thornton, J., Teh, Y. W., & Doucet, A. (2022). Riemannian score-based generative modeling. arXiv preprint arXiv:2202.02763.
[D]: Fishman, N., Klarner, L., De Bortoli, V., Mathieu, E., & Hutchinson, M. (2023). Diffusion Models for Constrained Domains. arXiv preprint arXiv:2304.05364.
[E]: Huang, C. W., Aghajohari, M., Bose, J., Panangaden, P., & Courville, A. C. (2022). Riemannian diffusion models. Advances in Neural Information Processing Systems, 35, 2750-2761.
[F]: Liu, L., Ren, Y., Lin, Z., & Zhao, Z. (2022). Pseudo numerical methods for diffusion models on manifolds. arXiv preprint arXiv:2202.09778.