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

1. Problem Statement

UNet architectures are ubiquitous in image restoration, segmentation, and diffusion models due to their simplicity and effectiveness. However, a critical theoretical limitation exists: while UNets are often assumed to be translation-equivariant (meaning shifting the input image results in an equivalent shift in the output), standard implementations are not.

The lack of true equivariance stems from aliasing introduced by several standard components:

• Pooling and Upsampling: Traditional max-pooling and zero-insertion upsampling introduce high-frequency artifacts.
• Activation Functions: Non-smooth functions like ReLU generate high-frequency components that alias during downsampling.
• Normalization: Standard Layer Normalization and Batch Normalization are not strictly translation-equivariant.
• Padding: Mismatched padding modes (e.g., zero-padding) create boundary errors that break equivariance.

Existing solutions have addressed these issues partially (e.g., in classifiers or generators) but have not been comprehensively applied to encoder-decoder UNet architectures, which require handling both downsampling and upsampling simultaneously.

2. Methodology: UNet-AF Architecture

The authors propose UNet-AF, a new UNet architecture designed from the ground up to be alias-free and translation-equivariant. The architecture is based on the UNet-Jin (a variant of the original UNet-Ronneberger with residual connections) but replaces every non-equivariant component with state-of-the-art alternatives.

Key architectural modifications include:

• Convolutions:
  • Circular Convolutions: Used for circular translation tasks to avoid boundary errors.
  • Zero-Padded Convolutions: Used for crop-translation tasks where information is discarded at boundaries.
• Pooling Layers (Downsampling):
  • Replaced Max-Pooling with BlurPool using periodic sinc anti-aliasing filters implemented via Fast Fourier Transform (FFT). This ensures every input pixel contributes to the output, eliminating aliasing within numerical precision.
• Upsampling Layers:
  • Replaced standard zero-insertion upsampling with Filtered Upsampling. This involves inserting zeros, applying a low-pass filter to remove aliasing, and then convolving.
• Activation Functions:
  • Replaced ReLU with Filtered GELU.
  • Mechanism: The input is resampled to a finer grid, the base activation (GELU) is applied, high frequencies are removed via a low-pass filter, and the result is resampled back to the coarse grid.
  • Rationale: GELU is chosen for its smoothness (limiting high-frequency generation) and linear growth (preventing amplification of numerical errors), unlike polynomial activations.
• Normalization Layers:
  • Replaced Batch Normalization with Alias-Free Layer Normalization (LayerNorm-AF).
  • Standard LayerNorm is known to be non-equivariant; the proposed LayerNorm-AF ensures equivariance while maintaining training stability.
• Residual Connections:
  • Made optional. The authors found that while they do not introduce aliasing, they are not always beneficial for stability or performance in this specific context.

3. Key Contributions

1. UNet-AF Architecture: The first comprehensive alias-free UNet designed for continuous sub-pixel translation equivariance, integrating anti-aliasing filters into every stage of the encoder-decoder pipeline.
2. Comprehensive Evaluation: Evaluated on multiple image restoration tasks (circular deblurring, valid deblurring, and denoising) comparing performance, training stability, and equivariance against standard baselines (UNet-Ronneberger and UNet-Jin).
3. Extensive Ablation Studies: Systematically validated that every proposed change (filtering, pooling, activation, normalization) is crucial. Removing any single alias-free component causes a drastic drop in equivariance (approx. 40 dB).
4. Trade-off Analysis: Demonstrated that while equivariance improves performance and stability, it comes at a significant computational cost (approx. 7x slower inference).

4. Results

The experiments were conducted on the DIV2K dataset with tasks including Gaussian deblurring and denoising.

• Equivariance (EQUIV Metric):
  • UNet-AF achieved an EQUIV score of 81.63 dB, compared to 37.60 dB for UNet-Ronneberger and 38.20 dB for UNet-Jin.
  • This represents a massive improvement, indicating the network output shifts almost perfectly with the input.
  • In ablation studies, substituting any alias-free layer with a standard one caused the EQUIV score to drop by roughly 40 dB, proving the necessity of the full suite of modifications.
• Performance (PSNR/SSIM/LPIPS):
  • UNet-AF achieved competitive or superior performance in PSNR and SSIM compared to baselines.
  • Notably, the use of Alias-Free LayerNorm significantly improved both performance and training stability.
  • Filtered GELU showed slightly better equivariance than Filtered ReLU, though polynomial activations resulted in poor performance.
• Training Stability:
  • UNet-AF exhibited much smoother training dynamics (lower variance in PSNR increments between epochs) compared to baselines.
  • Baselines showed erratic performance, whereas UNet-AF converged more reliably.
• Computational Cost:
  • The primary drawback is efficiency. UNet-AF runs at ~290 FPS compared to ~2494 FPS for the baseline UNet-Ronneberger (a ~7x slowdown).
  • The slowdown is attributed to the FFT-based anti-aliasing filters and the additional operations in normalization and activation layers.

5. Significance and Conclusion

The paper establishes that translation equivariance is not just a theoretical property but a practical driver for stability and performance in image restoration. By rigorously eliminating aliasing sources, UNet-AF achieves a level of robustness where the network's output is invariant to the position of features in the input image.

Key Takeaways:

• Necessity of Holistic Design: Equivariance cannot be achieved by fixing a single layer; the entire pipeline (pooling, upsampling, activation, normalization) must be alias-free.
• Robustness: The proposed model is significantly more robust to adversarial translations and sub-pixel shifts than standard UNets.
• Future Work: The authors suggest that while the current implementation is computationally expensive, future optimizations using dedicated GPU kernels and layer-specific anti-aliasing filters could mitigate the speed penalty.

In summary, UNet-AF sets a new standard for translation-equivariant deep learning in image restoration, proving that the theoretical requirement of equivariance yields tangible benefits in model stability and reconstruction quality, provided the computational cost is acceptable.