youssefsaied/normalization_equivariance

This repository is the official code release for a research paper accepted at ICML 2026, the International Conference on Machine Learning. It addresses a specific limitation in AI image denoising, the process of removing noise (graininess or static) from images. Most denoising models are trained at a fixed noise level. When applied to an image with a different amount of noise, a situation called noise-level mismatch, quality can drop sharply. This repo introduces WNE (Wrapped Normalization Equivariance), a technique that wraps around any existing denoising model to make it robust to this mismatch, without modifying the model's internal structure and with no measurable extra GPU overhead. The core mechanism works in three steps: before passing an image into the denoiser, WNE computes a global mean and standard deviation, normalizes the image into a consistent scale, runs it through the original model, then rescales the output back. This normalize-process-denormalize pattern is proven to be the exact characterization of normalization-equivariant behavior, meaning it is a mathematically rigorous solution, not a heuristic workaround. The repo includes training and evaluation code for several denoising architectures covered in the paper, DnCNN, FDnCNN, SwinIR, and Restormer, covering both CNN-based and transformer-based approaches. Evaluation scripts measure quality using standard metrics, and example commands reproduce the paper's figures and benchmark results. Training datasets and pretrained checkpoints are not bundled in the repository and must be obtained separately. The code is written in Python and released under an MIT license.