Normalization Equivariance for Arbitrary Backbones, with Application to Image Denoising

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.

AI systems that work on images often need to respond predictably when the brightness and contrast of a photo change. Previously, the only way to guarantee this was to rebuild the neural network from the inside out with specially designed parts — slow, and incompatible with the strongest modern networks, such as transformers. We show that a much simpler recipe achieves this same predictable response: strip away the image's overall brightness and contrast before the network sees it, run the network unchanged, and put the original brightness and contrast back at the end. The network itself is left untouched, and the math guarantees the behavior we want. Any network can now be given this predictability, including transformer-based ones previously out of reach, at no extra computing cost. Our main results show this makes image-cleaning networks far more reliable when faced with unfamiliar noise levels, and early results suggest the same idea helps with related tasks such as image reconstruction.