UNet-AF: An alias-free UNet for image restoration

TL;DR

This paper introduces UNet-AF, an alias-free UNet architecture that enhances translation equivariance in image restoration tasks by carefully selecting layers, leading to improved theoretical properties and competitive empirical performance.

Contribution

The paper presents a novel alias-free UNet architecture that significantly improves translation equivariance in image restoration, addressing aliasing issues in traditional UNets.

Findings

Enhanced translation equivariance in UNet-AF

Competitive image restoration performance

Each architectural change is crucial for equivariance

Abstract

The simplicity and effectiveness of the UNet architecture makes it ubiquitous in image restoration, image segmentation, and diffusion models. They are often assumed to be equivariant to translations, yet they traditionally consist of layers that are known to be prone to aliasing, which hinders their equivariance in practice. To overcome this limitation, we propose a new alias-free UNet designed from a careful selection of state-of-the-art translation-equivariant layers. We evaluate the proposed equivariant architecture against non-equivariant baselines on image restoration tasks and observe competitive performance with a significant increase in measured equivariance. Through extensive ablation studies, we also demonstrate that each change is crucial for its empirical equivariance. Our implementation is available at https://github.com/jscanvic/UNet-AF