Design choices of guidance methods

[nilsleh]

Hi @francois-rozet ,

Thanks for this fantastic (and greatly time-saving) package and also all your super interesting contributions (in particular data score based data assimiliation, which is a research area I am trying to get into)!

I had a question about the design choice of the guidance modules, namely that many conditional guidance schemes like MMPS are expected to be initialized with the observation y as an argument, as opposed to the observation being an argument to the __call__ function since this would allow calling the guidance module repeatedly with different observations, without instantiating a new module every time. So I was just wondering about the design choice to better understand the mechanics of all the different modules.

Cheers

[francois-rozet]

Hi @nilsleh, thank you for the kind words!

That's a great question, and I have actually been wondering whether I should modify this design choice. So, first of all, why is it currently like this?

For some guidance methods, it is possible to pre-compute some quantities based on $A$ and $y$ (and others). In that case, it is more efficient to instantiate the conditional denoiser with $A$ and $y$ once, and reuse the computed quantities during the denoiser's forward. It is the case of SILO, for example.

In addition, some pre-trained models (e.g. ADM and EDM) use y as the class label keyword. If A and y (and others) are provided in the forward, it is not possible to pass arguments with key y to the underlying denoiser.

Now, why would it be better to pass A and y in the forward? As you said, you would not need to re-instantiate a conditional denoiser for every observation (which isn't costly as weights are shared). That's more or less it. So I came to the conclusion that it was not worth the hassle.

Do you have a specific use case in mind where it would make a bigger difference?

[nilsleh]

Thanks for the fast reply that makes a lot of sense. I guess I just found it a bit unintuitive at first sight to be re-instantiating classes for every sample or batch of data that one would iterate over.

I was looking at the results of the MMPS paper in Table 4 of the Appendix regarding the results for inverse problems such as inpainting and super-resolution and tried to set this up with azula modules and code.

I was also looking in the MMPS benchmark repo, where several operators are defined: https://github.com/francois-rozet/mmps-benchmark/blob/877f659d6518e51a1a1f2dff1d5dfcc3b86e1c0d/guided_diffusion/measurements.py#L139

Since the call of the operator only accepts the single data argument

azula/azula/guidance/mmps.py

Line 79 in 6e5b8f3

y_hat = self.A(x_hat)

and we reinstantiate the sampler anyway, we would also just reinstante an operator that takes an init argument of the mask according to the data that will be passed.

So thanks again, I have a better understanding of the intended use of the modules.

[francois-rozet]

If you have a measurement A that takes other arguments than x, you can wrap it with functools.partial.

cond_denoiser = MMPSDenoiser(
    denoiser,
    y=y.to(device),
    A=functools.partial(A, mask=mask.to(device)),
)

Note that in Azula 0.5.0, the API has changed a little, with a few breaking changes. Notably, noise schedules and samplers are not modules anymore.