The DDPM U-Net assembles encoder, bottleneck, and decoder stages using ModuleList, with skip connections via concatenation

The full Conditional U-Net is assembled by iterating over resolution levels and building ModuleLists of [ResNetBlock, ResNetBlock, LinearAttention, Downsample/Conv] for the encoder and symmetrically [ResNetBlock, ResNetBlock, LinearAttention, Upsample/Conv] for the decoder. Skip connections are stored in a list h during the downward pass and concatenated via torch.cat during the upward pass. At the bottleneck, two ResNet blocks flank a full attention layer. This pattern — symmetrical encoder/decoder with skip-concatenated feature maps — is the definitive U-Net blueprint. The pattern generalises beyond DDPM to segmentation and other dense prediction tasks.

Examples

Encoder loop: block1(x,t) -> h.append(x) -> block2(x,t) -> attn(x) -> h.append(x) -> downsample(x). Decoder: cat(x, h.pop()) -> block1 -> cat(x, h.pop()) -> block2 -> attn -> upsample.

Assessment

Trace the U-Net forward pass for 3 resolution levels: list the operations in order and identify where each skip connection is saved and consumed.