# A Unified Framework for U-Net Design and Analysis

arXiv: https://arxiv.org/abs/2305.19638 · HTML: https://mlsanity.com/paper/2305.19638/

## Key facts

- Cited by 27 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 44 of its 99 references that are in the map.
- First posted on arXiv 31 May 2023.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 30 of 44)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding](https://mlsanity.com/paper/2205.11487/) (2022)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Improved Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2102.09672/) (2021)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [UNet++: A Nested U-Net Architecture for Medical Image Segmentation](https://mlsanity.com/paper/1807.10165/) (2018)
- [3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation](https://mlsanity.com/paper/1606.06650/) (2016)
- [Attention U-Net: Learning Where to Look for the Pancreas](https://mlsanity.com/paper/1804.03999/) (2018)
- [Variational Diffusion Models](https://mlsanity.com/paper/2107.00630/) (2021)
- [Diffusion Schr"odinger Bridge with Applications to Score-Based Generative Modeling](https://mlsanity.com/paper/2106.01357/) (2021)
- [Graph U-Nets](https://mlsanity.com/paper/1905.05178/) (2019)
- [Reliable Fidelity and Diversity Metrics for Generative Models](https://mlsanity.com/paper/2002.09797/) (2020)
- [Simple diffusion: End-to-end diffusion for high resolution images](https://mlsanity.com/paper/2301.11093/) (2023)
- [PDEBENCH: An Extensive Benchmark for Scientific Machine Learning](https://mlsanity.com/paper/2210.07182/) (2022)
- [Assessing Generative Models via Precision and Recall](https://mlsanity.com/paper/1806.00035/) (2018)
- [Invertible Residual Networks](https://mlsanity.com/paper/1811.00995/) (2018)
- [Road Extraction by Deep Residual U-Net](https://mlsanity.com/paper/1711.10684/) (2017)
- [nnU-Net: Self-adapting Framework for U-Net-Based Medical Image Segmentation](https://mlsanity.com/paper/1809.10486/) (2018)
- [A Probabilistic U-Net for Segmentation of Ambiguous Images](https://mlsanity.com/paper/1806.05034/) (2018)
- [Very Deep VAEs Generalize Autoregressive Models and Can Outperform Them on Images](https://mlsanity.com/paper/2011.10650/) (2020)
- [Riemannian Score-Based Generative Modelling](https://mlsanity.com/paper/2202.02763/) (2022)
- [Towards Multi-spatiotemporal-scale Generalized PDE Modeling](https://mlsanity.com/paper/2209.15616/) (2022)

## Built on by (most cited first, 27 of 27)

- [U-DiTs: Downsample Tokens in U-Shaped Diffusion Transformers](https://mlsanity.com/paper/2405.02730/) (2024)
- [A Fourier Space Perspective on Diffusion Models](https://mlsanity.com/paper/2505.11278/) (2025)
- [A Survey of Mamba](https://mlsanity.com/paper/2408.01129/) (2024)
- [U-REPA: Aligning Diffusion U-Nets to ViTs](https://mlsanity.com/paper/2503.18414/) (2025)
- [Effective Diffusion Transformer Architecture for Image Super-Resolution](https://mlsanity.com/paper/2409.19589/) (2024)
- [Spectral Motion Alignment for Video Motion Transfer using Diffusion Models](https://mlsanity.com/paper/2403.15249/) (2024)
- [T2ICount: Enhancing Cross-modal Understanding for Zero-Shot Counting](https://mlsanity.com/paper/2502.20625/) (2025)
- [U-Nets as Belief Propagation: Efficient Classification, Denoising, and Diffusion in Generative Hierarchical Models](https://mlsanity.com/paper/2404.18444/) (2024)
- [The Surprising Effectiveness of Skip-Tuning in Diffusion Sampling](https://mlsanity.com/paper/2402.15170/) (2024)
- [PixelU: A U-Shaped Transformer for Efficient End-to-End Pixel Diffusion](https://mlsanity.com/paper/2606.27760/) (2026)
- [Neural Network Emulator for Atmospheric Chemical ODE](https://mlsanity.com/paper/2408.01829/) (2024)
- [Rethinking Iterative Stereo Matching from Diffusion Bridge Model Perspective](https://mlsanity.com/paper/2404.09051/) (2024)
- [Demystifying the Effect of Receptive Field Size in U-Net Models for Medical Image Segmentation](https://mlsanity.com/paper/2406.16701/) (2024)
- [Responsible Diffusion: A Comprehensive Survey on Safety, Ethics, and Trust in Diffusion Models](https://mlsanity.com/paper/2509.22723/) (2025)
- [Score-based generative emulation of impact-relevant Earth system model outputs](https://mlsanity.com/paper/2510.04358/) (2025)
- [TexTailor: Inference-Time Textual Guidance Tailoring for Multimodal Diffusion Transformers](https://mlsanity.com/paper/2601.02211/) (2026)
- [Fourier Neural Operators for Learning Dynamics in Quantum Spin Systems](https://mlsanity.com/paper/2409.03302/) (2024)
- [It's Not the Target, It's the Background: Rethinking Infrared Small Target Detection via Deep Patch-Free Low-Rank Representations](https://mlsanity.com/paper/2506.10425/) (2025)
- [QID2: An Image-Conditioned Diffusion Model for Q-space Up-sampling of DWI Data](https://mlsanity.com/paper/2409.02309/) (2024)
- [CHEM: Estimating and Understanding Hallucinations in Deep Learning for Image Processing](https://mlsanity.com/paper/2512.09806/) (2025)
- [Object-Size-Driven Design of Convolutional Neural Networks: Virtual Axle Detection based on Raw Data](https://mlsanity.com/paper/2309.01574/) (2023)
- [Towards Robust Unsupervised Attention Prediction in Autonomous Driving](https://mlsanity.com/paper/2501.15045/) (2025)
- [Bringing together invertible UNets with invertible attention modules for memory-efficient diffusion models](https://mlsanity.com/paper/2504.10883/) (2025)
- [Seeing Through the Brain: New Insights from Decoding Visual Stimuli with fMRI](https://mlsanity.com/paper/2510.16196/) (2025)
- [Multi-ResNets for Subspace Preconditioning in Constrained Optimization](https://mlsanity.com/paper/2606.06300/) (2026)
- [U-shaped Multi-granularity Learning for Vision-Language Models](https://mlsanity.com/paper/2607.14966/) (2026)
- [Hybrid Lagrangian-Eulerian Model for Lagrangian Fluid Simulation](https://mlsanity.com/paper/2608.01164/) (2026)

## Abstract (arXiv)

U-Nets are a go-to, state-of-the-art neural architecture across numerous tasks for continuous signals on a square such as images and Partial Differential Equations (PDE), however their design and architecture is understudied. In this paper, we provide a framework for designing and analysing general U-Net architectures. We present theoretical results which characterise the role of the encoder and decoder in a U-Net, their high-resolution scaling limits and their conjugacy to ResNets via preconditioning. We propose Multi-ResNets, U-Nets with a simplified, wavelet-based encoder without learnable parameters. Further, we show how to design novel U-Net architectures which encode function constraints, natural bases, or the geometry of the data. In diffusion models, our framework enables us to identify that high-frequency information is dominated by noise exponentially faster, and show how U-Nets with average pooling exploit this. In our experiments, we demonstrate how Multi-ResNets achieve competitive and often superior performance compared to classical U-Nets in image segmentation, PDE surrogate modelling, and generative modelling with diffusion models. Our U-Net framework paves the way to study the theoretical properties of U-Nets and design natural, scalable neural architectures for a multitude of problems beyond the square.
