Total Variation Regularization

The problem, from first principles

Optimizing the pixels of an image to match feature statistics tends to produce high-frequency noise: speckles that fool the network's features but look bad to a person. Real images are mostly piecewise smooth, with neighbouring pixels similar except at edges. Total variation (TV) penalizes the differences between neighbouring pixels, so adding it to the loss acts as a prior toward clean images. The L1 form used here punishes noise but still allows sharp edges to survive better than a squared penalty, which would blur them. It is the same idea as a denoising prior and is a standard regularizer for any pixel-space optimization.

From theory to code

Implement total_variation.

Constraints

• img has shape (C, H, W) or (H, W); always work on the last two axes.
• Return 0.5 * (mean(|img[..., 1:, :] - img[..., :-1, :]|) + mean(|img[..., :, 1:] - img[..., :, :-1]|)) as a float: the average absolute vertical difference plus the average absolute horizontal difference, halved.

Hints