Normalization Equivariance for Arbitrary Backbones, with Application to Image Denoising

TL;DR

This paper introduces Wrapped Normalization Equivariance (WNE), a parameter-free method that enhances robustness of CNNs and transformers in image denoising without additional GPU cost.

Contribution

The paper proposes WNE, a novel wrapper that enforces normalization equivariance efficiently, applicable to arbitrary backbones, and proves its theoretical equivalence to NE functions.

Findings

WNE improves robustness to noise-level mismatch in denoising tasks.

WNE adds no measurable GPU overhead compared to baseline architectures.

Architectural NE baselines are up to 1.6x slower than WNE.

Abstract

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(a y + b\mathbf{1}) = a f(y) + b\mathbf{1}$ for all $a>0$ and $b\in\mathbb{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\times$ slower.