L2-norm discrepancy between paper and code

[tristan-deep]

Careful note for people implementing DPS!

The DPS paper specifies a squared L2 norm (line 7 of the algo):

However, in the code a regular L2 norm implemented with torch.linalg.norm is used:

diffusion-posterior-sampling/guided_diffusion/condition_methods.py

Line 31 in effbde7

norm = torch.linalg.norm(difference)

We've found DPS to work well with the non-squared version, as implemented in this codebase :) See our implementation amongst other diffusion guidance methods here.

[JeffFessler]

I came here to report the same issue and discovered that it had already been reported.
However, after discussing further with my collaborators, I now think it is deliberate and neither a bug nor a discrepancy!

Footnote 5 and Appendix D.1 of the 2023 ICLR paper on DPS explain that the authors recommend using a step size that is proportional to 1 / ‖ y - A x₀ ‖.
If you take the gradient with respect to the 2-norm squared and then multiply it by that reciprocal step size, you get the same result as taking the gradient with respect to the non-squared version, within a factor of 2.

In fact, in line 32, immediately below line 31 that you quote above, the variable name is norm_grad

diffusion-posterior-sampling/guided_diffusion/condition_methods.py

Line 32 in effbde7

norm_grad = torch.autograd.grad(outputs=norm, inputs=x_prev)[0]

which most likely means "normalized gradient," i.e., the gradient of the 2-norm squared function, normalized by the 2-norm of the residual, which is proportional to the gradient of the non-squared version.

This is a subtle point, and it sure would have been nice if the authors had left comments in the code to help readers appreciate the intention. Sadly, uncommented code seems quite common, despite the promises of "open source" and "reproducible research." I also wish they had included the usual factor of 1/2 in line 7 of the pseudo-code above, because that would have eliminated the factor-of-2 discrepancy between the code and the pseudo-code.

[tristan-deep]

Sneaky! Well spotted :) And then it seems to make sense indeed!

Also agree, this nuance is a bit hidden. Both in the code as well as in the paper. @hyungjin-chung maybe can add a small comment around here to highlight this fact?