LION.losses.jacobian_penalty

Penalties on a denoiser’s Jacobian, for Plug-and-Play methods that converge.

PnP convergence proofs ask something of the denoiser D: that it be firmly nonexpansive (2D - I nonexpansive; Pesquet et al. 2021), or that its residual I - D, the gradient of the potential of a gradient-step denoiser, be L-Lipschitz with L < 1 (Hurault et al. 2022). A network isn’t, so it is trained with a penalty on the spectral norm of the Jacobian of that map, estimated by power iterations. They are model regularisations: a solver adds one to its loss when it is its model_regularization (off by default), called as penalty(model, step) with the mini-batch’s step (LIONsolver).

Classes

FirmlyNonexpansivePenalty([weight, epsilon, ...])	Pesquet et al. (2021): weight * max(||J_Q(x̃)||², 1 - epsilon), with Q = 2D - I and x̃ a random point between each clean image and its denoised image (x̃ = r x + (1 - r) D(y), r uniform in [0, 1]).
ProximalDenoiserPenalty([weight, epsilon, ...])	Hurault et al. (2022): weight * max(||J_{I-D}(y)||, 1 - epsilon) at the noisy images y: the Hessian of a gradient-step denoiser's potential (I - D = ∇g), so that D is a proximal operator.

Functions

jacobian_spectral_norm(f, x[, power_iterations])

The spectral norm of the Jacobian J of f at each image of the batch x (f maps each image on its own), estimated with power_iterations iterations of the power method on JᵀJ.

---

class LION.losses.jacobian_penalty.FirmlyNonexpansivePenalty(weight=1e-05, epsilon=0.05, power_iterations=50)

Bases: Module

Pesquet et al. (2021): weight * max(||J_Q(x̃)||², 1 - epsilon), with Q = 2D - I and x̃ a random point between each clean image and its denoised image (x̃ = r x + (1 - r) D(y), r uniform in [0, 1]).

Their settings: weight 1e-5, epsilon 0.05, after training without it (pretraining). They estimate the norm with the power method (their Appendix C); 50 iterations is Hurault et al.’s number.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(model, step)

The penalty for a mini-batch step: the denoiser step.forward, its clean images step.target, its outputs step.output.

static cite(cite_format='bib')

Where the penalty comes from: its citation, BibTeX (“bib”) or “MLA”.

class LION.losses.jacobian_penalty.ProximalDenoiserPenalty(weight=0.001, epsilon=0.1, power_iterations=50)

Bases: Module

Hurault et al. (2022): weight * max(||J_{I-D}(y)||, 1 - epsilon) at the noisy images y: the Hessian of a gradient-step denoiser’s potential (I - D = ∇g), so that D is a proximal operator.

Their settings: weight (μ) 1e-3 or 1e-2, epsilon 0.1, 50 power iterations, fine-tuning a trained gradient-step denoiser for 10 epochs.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(model, step)

The penalty for a mini-batch step: the denoiser step.forward at its inputs step.input.

static cite(cite_format='bib')

Where the penalty comes from: its citation, BibTeX (“bib”) or “MLA”.

LION.losses.jacobian_penalty.jacobian_spectral_norm(f, x, power_iterations=50)

The spectral norm of the Jacobian J of f at each image of the batch x (f maps each image on its own), estimated with power_iterations iterations of the power method on JᵀJ. Differentiable with respect to f’s parameters (as ||Jᵀv||, v the last left singular vector: second derivatives only, which e.g. batch normalisation supports).