sbi icon indicating copy to clipboard operation
sbi copied to clipboard

Adding interface to azula samplers

Open gmoss13 opened this issue 9 months ago • 0 comments

🚀 Feature Request

Is your request related to a problem?

Currently, we only have a few diffusion samplers for trained score estimators (see #1438). In addition to implementing more options, it would be nice to add an interface for samplers from azula. Azula has a lot of diffusion samplers already implemented, potentially saving us a lot of work. Additionally, they are designed to be plug-and-play, so we should be able to pass our trained score estimators to azula samplers without too much interface engineering.

Describe the solution you'd like

Typically, we create a diffusion sampler object by specifiying a potential function, a predictor, and a corrector:

diffuser = Diffuser(
    self.potential_fn,
    predictor=predictor,
    corrector=corrector,
    predictor_params=predictor_params,
    corrector_params=corrector_params,
)

azula handles samplers a bit differently, as the different samplers already correspond to different predictor-corrector combinations. An easier solution would be to just wrap azula samplers as Predictor classes, and forcing corrector=None if any azula-based sampler is used for Predictor. A more robust solution is to write an AzulaDiffuser class that wraps around azula samplers directly and inherits from the base Diffuser class, where the predictor and corrector parameters are passed on to the azula samplers.

gmoss13 avatar Mar 18 '25 09:03 gmoss13