# U-Net vs Diffusion UNet

A segmentation U-Net against the one a diffusion model uses.

**Diffusion UNet has 6.0M more parameters than U-Net: 7 layers added, 12 removed, 11 changed.**

Source: https://neurarch.com/diff/unet-vs-diffusion-unet.html

## Sides

| | U-Net | Diffusion UNet |
|---|---|---|
| Layers | 22 | 17 |
| Parameters | 721K | 6.7M |
| Input | 3 × 256 × 256 | 4 × 64 × 64 |
| Output | 1 × 256 × 768 | 4 × 64 × 64 |
| Forward-passes | yes | yes |
| Est. train cost | $0.593 | $0.541 |
| T4 16GB | fits | fits |
| A100 40GB | fits | fits |
| H100 80GB | fits | fits |

## Deltas (Diffusion UNet relative to U-Net)

- Parameters: +6.0M (9.3× the size)
- Layers: -5
- Added 7, removed 12, changed 11, unchanged 1

## Layer by layer

| # | Status | U-Net | Params | Output | Diffusion UNet | Params | Output |
|---|---|---|---|---|---|---|---|
| 1 | changed (shape) | image (Input) |  | 3 × 256 × 256 | noisy_latent (Input) |  | 4 × 64 × 64 |
| 2 | removed | enc1_conv (Conv2d) | 640 | 64 × 256 × 256 |  | |  |
| 3 | removed | enc1_bn (Batch Norm) | 128 | 64 × 256 × 256 |  | |  |
| 4 | removed | enc1_relu (Relu) |  | 64 × 256 × 256 |  | |  |
| 5 | removed | enc1_pool (Maxpool2d) |  | 64 × 128 × 128 |  | |  |
| 6 | changed (outChannels) | enc2_conv (Conv2d) | 1.3K | 128 × 128 × 128 | conv_in (Conv2d) | 3.2K | 320 × 64 × 64 |
| 7 | changed (type, numFeatures, numGroups, numChannels) | enc2_bn (Batch Norm) | 256 | 128 × 128 × 128 | down1_norm (Group Norm) | 640 | 320 × 64 × 64 |
| 8 | removed | enc2_relu (Relu) |  | 128 × 128 × 128 |  | |  |
| 9 | removed | enc2_pool (Maxpool2d) |  | 128 × 64 × 64 |  | |  |
| 10 | changed (outChannels) | bottleneck_conv (Conv2d) | 2.6K | 256 × 64 × 64 | down1_conv (Conv2d) | 3.2K | 320 × 64 × 64 |
| 11 | removed | bottleneck_bn (Batch Norm) | 512 | 256 × 64 × 64 |  | |  |
| 12 | changed (type) | bottleneck_relu (Relu) |  | 256 × 64 × 64 | down1_silu (Swish) |  | 320 × 64 × 64 |
| 13 | removed | up2 (Transpose Conv2d) | 640 | 128 × 128 × 128 |  | |  |
| 14 | removed | dec2_skip (Concatenate) |  | 128 × 128 × 256 |  | |  |
| 15 | added |  | |  | to_tokens (Reshape) |  | 4096 × 320 |
| 16 | added |  | |  | down1_text_attn (Cross Attention) | 411K | 4096 × 320 |
| 17 | added |  | |  | to_feature_map (Reshape) |  | 320 × 64 × 64 |
| 18 | changed (outChannels, stride) | dec2_conv (Conv2d) | 1.3K | 128 × 128 × 256 | downsample_1 (Conv2d) | 6.4K | 640 × 32 × 32 |
| 19 | changed (type, numFeatures, numGroups, numChannels) | dec2_bn (Batch Norm) | 256 | 128 × 128 × 256 | mid_norm (Group Norm) | 1.3K | 640 × 32 × 32 |
| 20 | removed | dec2_relu (Relu) |  | 128 × 128 × 256 |  | |  |
| 21 | removed | up1 (Transpose Conv2d) | 320 | 64 × 256 × 512 |  | |  |
| 22 | removed | dec1_skip (Concatenate) |  | 64 × 256 × 768 |  | |  |
| 23 | added |  | |  | to_tokens (Reshape) |  | 1024 × 640 |
| 24 | added |  | |  | mid_text_attn (Cross Attention) | 1.6M | 1024 × 640 |
| 25 | added |  | |  | to_feature_map (Reshape) |  | 640 × 32 × 32 |
| 26 | added |  | |  | upsample_1 (Upsample) |  | 640 × 64 × 64 |
| 27 | changed (outChannels) | dec1_conv (Conv2d) | 640 | 64 × 256 × 768 | up1_conv (Conv2d) | 3.2K | 320 × 64 × 64 |
| 28 | changed (type, numFeatures, numGroups, numChannels) | dec1_bn (Batch Norm) | 128 | 64 × 256 × 768 | conv_out_norm (Group Norm) | 640 | 320 × 64 × 64 |
| 29 | changed (type) | dec1_relu (Relu) |  | 64 × 256 × 768 | up1_silu (Swish) |  | 320 × 64 × 64 |
| 30 | changed (outChannels, kernelSize, padding) | output_conv (Conv2d) | 2 | 1 × 256 × 768 | conv_out (Conv2d) | 40 | 4 × 64 × 64 |
| 31 | same | segmentation_mask (Output) |  | 1 × 256 × 768 | predicted_noise (Output) |  | 4 × 64 × 64 |

## What this is not

- The two are priced at different declared inputs (3 × 256 × 256 against 4 × 64 × 64), so memory, cost and GPU fit are each right about their own model and are not a comparison between them. The layer and parameter deltas are unaffected.
- Parameter counts are derived from the graph, not read from a checkpoint. They are exact for a graph that is fully specified and approximate for one that is not.
- Cost and GPU fit are estimates from the graph under one set of assumptions, not measurements of a run.

## Graphs

- U-Net: https://neurarch.com/templates/unet/model.json
- Diffusion UNet: https://neurarch.com/templates/diffusion-unet/model.json
- Check a graph of your own: `POST https://www.neurarch.com/api/v1/plan`