The simplicity and effectiveness of UNet architectures make them ubiquitous in image restoration, 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 show how to build sub-pixel translation-equivariant UNet architectures by appropriately choosing their main components (convolution, pooling, downsampling, activation, and normalization layers) to be alias-free. We evaluate the proposed equivariant architectures 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 the importance of every architectural choice to achieve high equivariance. Our implementation is available at https://github.com/jscanvic/UNet-AF
Figures & tables
Fig. 1: Equivariance and robustness to circular translations. (Top) Performance on an horizontally translated image, with 0.01 px steps, (bottom) adversarial performance on the worst translation up to a maximum displacement, with 0.25 px steps. UNet-AF is stable and robust, unlike the baselines.
Fig. 2: Overview of the architecture. UNet-AF is obtained by changing the architectures of [ 1 ] , and [ 2 ] . Anti-aliasing filters are added in upsampling layers, max-pooling layers are replaced with blur pooling layers, ReLU activations are replaced with filtered GELU layers and batch normalization is replaced by alias-free layer normalization. The padding mode and residual connection are specified on a per-task basis.
Method
Difference
PSNR ↑
SSIM ↑
LPIPS ↓
EQUIV ↑
STAB ↑
FPS ↑
UNet-Ronneberger
31.28 ± 2.13
0.9316 ± 0.0271
0.0377 ± 0.0235
37.60 ± 2.62
9.10
2494
UNet-Jin
+ Res conn.
31.64 ± 2.51
0.9280 ± 0.0345
0.0402 ± 0.0254
38.20 ± 2.22
12.86
2114
+ LayerNorm-AF
32.16 ± 2.37
0.9383 ± 0.0211
0.0355 ± 0.0210
38.41 ± 1.99
18.40
1230
+ Circular pad.
32.46 ± 2.45
0.9417 ± 0.0209
0.0323 ± 0.0199
43.16 ± 2.38
20.63
982
+ Filtered GELU
32.22 ± 2.39
0.9383 ± 0.0201
0.0331 ± 0.0193
43.55 ± 2.77
20.48
302
+ Filtered Upconv.
32.24 ± 2.35
0.9382 ± 0.0201
0.0333 ± 0.0197
46.33 ± 2.59
20.34
293
TABLE I: Performance on circular deblurring. The changes in Section II lead to a steady increase in equivariance (EQUIV), stabilize the training dynamics and improve performance (PSNR, SSIM, LPIPS), at the expense of a higher computational cost (FPS). Best metric in bold , second best underlined . Values: avg ± st.d.
Fig. 3: Evaluation PSNR throughout training for circular deblurring. Our proposed model has significantly smoother training dynamics than the non-equivariant baselines.
Replacement
Variant
PSNR ↑
EQUIV ↑
Res. connection
Present
32.22 ± 2.32
81.63 ± 0.26
Absent
31.96 ± 2.33
77.77 ± 2.32
Normalization
LayerNorm-AF
32.22 ± 2.32
81.63 ± 0.26
BatchNorm
32.21 ± 2.31
80.62 ± 8.14
InstanceNorm
31.50 ± 2.19
78.40 ± 0.42
LayerNorm
32.31 ± 2.36
54.19 ± 3.03
TABLE II: Ablation study. Each design choice made in Section II is explored in terms of performance and equivariance. Best metric in bold , second best underlined . Values: avg ± st.d.
Fig. 4: Reconstructions of a circularly blurred image. The PSNR is displayed in the bottom-left corner of each image.
Method
Task
PSNR ↑
EQUIV ↑
UNet-Ronneberger
Valid deblurring
31.67 ± 2.33
39.96 ± 2.71
UNet-Jin
32.25 ± 2.41
42.82 ± 2.96
UNet-AF
32.30 ± 2.42
46.57 ± 2.58
Blurred images
26.85 ± 2.55
N/A
UNet-Ronneberger
Gaussian denoising
29.70 ± 1.56
39.80 ± 1.31
UNet-Jin
29.58 ± 1.61
40.64 ± 1.78
TABLE III: Performance on additional tasks. We evaluate UNet-AF on valid deblurring, and Gaussian denoising. Best metric in bold , second best underlined . Values: avg ± st.d.
Filters
Padding
PSNR ↑
SSIM ↑
LPIPS ↓
EQUIV ↑
STAB ↑
FPS ↑
Kaiser
Zeros
32.41 ± 2.42
0.9404 ± 0.0209
0.0363 ± 0.0210
43.12 ± 1.89
22.30
556
Kaiser
Circular
26.79 ± 2.24
0.8315 ± 0.0367
0.2247 ± 0.0806
89.35 ± 2.35
13.45
571
Tuned
Zeros
32.45 ± 2.44
0.9413 ± 0.0211
0.0362 ± 0.0210
43.94 ± 2.06
23.29
534
Tuned
Circular
32.39 ± 2.40
0.9408 ± 0.0207
0.0307 ± 0.0181
49.36 ± 2.84
23.29
552
TABLE IV: Ablation on UNet-AF Fast. Best metric in bold , second best underlined . Values: avg ± st.d.
In this work, we propose Fast Equivariant Imaging (FEI), a novel unsupervised learning framework to rapidly and efficiently train deep imaging networks without ground-truth data. FEI reformulates the EI objective through an inexact variable-splitting scheme, decoupling network training from an auxiliary restoration step implemented with a plug-and-play denoiser, this novel unsupervised scheme shows superior efficiency and performance compared to the standard Equivariant Imaging paradigm. In particular, our FEI schemes achieve an order-of-magnitude (10x) acceleration over standard EI on training U-Net for X-ray CT reconstruction and image inpainting, with improved generalization performance. Beyond offline training, the proposed scheme also enables efficient test-time adaptation of a pretrained model to individual samples, to secure further performance improvements. Extensive experiments show that the proposed approach provides a noticeable efficiency and performance gain over existing unsupervised methods and model adaptation techniques.
Guixian Xu, Jinglai Li, Junqi Tang
School of Mathematics, University of Birmingham, Birmingham B15 2TT, United Kingdom
Image restoration is an inherently ill posed inverse problem. Equivariant networks that embed geometric symmetry priors can mitigate this ill posedness and improve performance. However, current understanding of the relationship between network equivariance and data symmetry remains largely heuristic. Particularly for real world data with imperfect symmetry, existing research lacks a systematic theoretical framework to quantify symmetry, select transformation groups, or evaluate model data alignment. To bridge this gap, we conduct an analysis from an optimization perspective and formalize the intrinsic relationship among data symmetry priors, model equivariance, and generalization capability. Specifically, we propose for the first time a quantifiable definition of non strict symmetry at the dataset level (rather than sample level) and use it as a constraint to formulate the restoration inverse problem. We then show that the equivariance for restoration models can be naturally derived from this inverse problems incorporated the proposed symmetry constraints, and that the equivariance error of the optimal restoration operator is strictly bounded by the data symmetry error and the discretization mesh size. Furthermore, by analyzing the network's empirical risk, we demonstrate that aligning equivariance with data symmetry optimizes the bias variance trade off, minimizing the total expected risk. Guided by these insights, we propose a Sample Adaptive Equivariant Network that uses a hypernetwork and transformation learnable equivariant convolutions to dynamically align with each sample's inherent symmetry. Extensive experiments on super resolution, denoising, and deraining validate our theoretical findings and show significant superiority over standard baselines and traditional equivariant models. Our code and supplementary material are available at https://github.com/tanfy929/SA-Conv.
Feiyu Tan, Qi Xie, Zongben Xu +1
School of Mathematics and Statistics, Xi’an Jiaotong University, Xi’an, Shaanxi, China
Normalization Equivariance (NE) is a structural prior that improves robustness to distribution shift in image-to-image tasks. A function f is normalization equivariant iff f(ay+b1)=af(y)+b1 for all a>0 and b∈R. Existing NE methods constrain every internal layer to NE-compatible operations. These constraints add runtime cost and exclude standard transformer components such as softmax attention and LayerNorm. We introduce Wrapped Normalization Equivariance (WNE), a parameter-free wrapper that normalizes the input, applies any backbone, and denormalizes the output. We prove every NE function admits this factorization, so the wrapper exactly parameterizes the class of NE functions. On blind denoising, wrapping CNN and transformer architectures improves robustness under noise-level mismatch with no measurable GPU overhead, while architectural NE baselines are up to 1.6× slower.
Youssef Saied, François Fleuret
Department of Computer Science, University of Geneva, Switzerland · Meta FAIR.