Comparisons / U-Net vs Diffusion UNet

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.

Baseline

U-Net

Layers: 22
Parameters: 721K
Input: 3 × 256 × 256
Output: 1 × 256 × 768
Forward-passes: yes
Est. train cost: $0.593

Compared

Diffusion UNet

Layers: 17
Parameters: 6.7M
Input: 4 × 64 × 64
Output: 4 × 64 × 64
Forward-passes: yes
Est. train cost: $0.541

The deltas

Every number is Diffusion UNet relative to U-Net.

Parameters

+6.0M (9.3× the size)

Layers

-5

Added

Removed

Changed

Unchanged

Which GPUs each one fits

Each side is measured at its own declared input (3 × 256 × 256 against 4 × 64 × 64). Both columns are right about their own model; the difference between them is not a fact about the designs.

GPU	U-Net	Diffusion UNet
T4 16GB	fits	fits
A100 40GB	fits	fits
H100 80GB	fits	fits

Layer by layer

Aligned in topological order. 1 of 31 rows are the same layer with the same parameters.

Hide all 31 rows

		U-Net			Diffusion UNet
		Layer	Params	Output	Layer	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.

Take it further

Open either graph in the editor, change it, and check it again: U-Net · Diffusion UNet

Compare any two models of your own, including anything on Hugging Face: the comparison tool.

Machine-readable: this page as markdown · the pair index · POST https://www.neurarch.com/api/v1/plan for a graph of your own.

Related comparisons

Diffusion UNet vs DiT-XL/2

Convolutional against transformer backbones for diffusion.