4
$\begingroup$

I want to know why diffusion models always use U-Net.

In my opinion, they use U-Net because you can see features of different resolutions and skip connection is good to add detail of images. But I am not sure if it is the main reason why they use U-Net.

Are there other reasons they choose U-Net rather than other architectures?

$\endgroup$

3 Answers 3

6
$\begingroup$

I don't have a definitive answer but I'd state my intuitions anyways:

Diffusion models are highly related to the idea of stacked denoising autoencoders [Kumar et al. (2014)]. Additionally, U-Net-like architectures are a very common architecture for autoencoders on images. Here, I would agree with your intuition that the bottleneck and the skip-connections help with denoising because they provide representations of different granularity.

Another thought is that U-Nets are empirically well-suited for image segmentation problems. Despite the fact that pixels are classified individually, you want the segmentation output to have consistent areas on top of objects in the image. So you kind of turn a very noisy segmentation mask (original image input) into a mask with much less noise (segmented output).

I think the latter is debatable, I'd be happy to hear your thoughts.

$\endgroup$
6
  • $\begingroup$ Hi, thanks for reply. The "bottleneck" you mean is 1x1 conv or the process of lowering resolution? If it is the former, I saw an article that says diffusion model cares about output more than latent representations which I think is true. Maybe this is also one of the reasons why they remove bottleneck layer from architecture? As for image segmentation, because I am not familiar with it, I am sorry that I cannot see the connection between them XD. $\endgroup$ Commented Sep 25, 2022 at 6:05
  • 2
    $\begingroup$ With bottleneck I mean the decrease in layer size. This forces the encoding to discard some information which should improve denoising properties. The noise cannot be encoded lower dimensions. $\endgroup$
    – Chillston
    Commented Sep 26, 2022 at 11:04
  • 2
    $\begingroup$ Regarding the image segmentation part: Just look at this image and you know what image segmentation means. What you see there is the model output (colors) superimposed on the model input (a regular image). The goal of the model is to classify every pixel based on which object it belongs to (colors denote object classes). Thus, you could think of this task as the transformation of a very noisy version of this segmentation map (which is the regular image) to a very clean segmentation map. I hope this made it somewhat clearer? $\endgroup$
    – Chillston
    Commented Sep 26, 2022 at 11:06
  • 1
    $\begingroup$ Yes, it is clearer, thanks for explanation! If that is the case, I think I prefer the former idea. Maybe it is really just a simple reason like this that supports them to use it :o. $\endgroup$ Commented Sep 26, 2022 at 11:19
  • $\begingroup$ I think stable diffusion uses UNet because diffusion probabilistic models used UNet, DPM used UNet because PixelCNN++ used UNet, PixelCNN++ used it because it was 2017, that's what they had and it was good enough. Deep learning research has always been about cargo cult empirical research with weak (if any) justification. That said, the other choice is stuff like MaskRCNN and RPNs don't make any sense on latent representations $\endgroup$
    – stuart
    Commented May 14, 2023 at 5:45
1
$\begingroup$

They do not.

Even "the" seminal paper does not (it was put on axiv before "U-Net" was a thing :).

$\endgroup$
0
$\begingroup$

Generally, U-Nets are the default choice when dealing with functions where inputs and outputs have the same size/shape (such as in image segmentation) due to the lack of a real information bottleneck. Thus, those are a natural good choice when dealing with diffusion models, where the network must predict the residual noise.

$\endgroup$

You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .