Diffusion models

[francois-rozet]

I am considering creating a small library for diffusion/score-based generative models. The idea would be to have a very simple, composable and extendable interface. For example, it would define the Denoiser, Scheduler and Sampler primitives (similar to the code in francois-rozet/diffusion-priors but in PyTorch).

I have several questions for Zuko's community:

1. Is there interest for such diffusion models library?
2. What would be the features you would like for such library? (e.g. conditioning methods, fidelity to papers vs standardized notations, pre-trained weights, ...)
3. Should this library be part of Zuko or should it be a separate library? In the former case, Zuko would become a "generative models" library instead of a normalizing flow library.

cc: @glouppe, @gerome-andry, @JuliaLinhart, @bkmi, @michaeldeistler, @janfb, @simonschnake

[bkmi]

1. I don't know if I'd use it personally so often, but I think the simplicity/style of zuko is really nice. I don't know of another library for diffusion that's so nice. Therefore I think there would be interest.

Furthermore, zuko essentially expanded the utilization of flows in sbi dramatically. I imagine such a library could also seperate the diffusion method from the research area in a positive way: e.g. more diffusion-based sbi algorithms.

2. I think lucid rains tries to copy papers fairly successfully. That makes me think it's better to have a framework which supports alternative options and then just implement those options in the framework--like zuko does.

3. Imo there is value in modularity, especially since there are dependencies on zuko now. As long as one can install both and use them together, that's the best imo. The question is: should they have some parent library that they both depend on? I would wait to see if it's necessary.

[francois-rozet]

Thank you Ben. So there is interest, but you think that a separate library would be more appropriate? This is my opinion as well, especially as adding diffusion models in Zuko could require significant API changes.

The question is: should they have some parent library that they both depend on? I would wait to see if it's necessary.

I think they can be independent and there is no real need for a parent library (except for torch).

That makes me think it's better to have a framework which supports alternative options and then just implement those options in the framework--like zuko does.

I am not sure to understand this answer. What I got is that you think that it is better to have a common framework (unified notations), which is expressive enough to implement special cases (paper quirks). Am I right?

[bkmi]

Thank you Ben. So there is interest, but you think that a separate library would be more appropriate? This is my opinion as well, especially as adding diffusion models in Zuko could require significant API changes.

agree!

I am not sure to understand this answer. What I got is that you think that it is better to have a common framework (unified notations), which is expressive enough to implement special cases (paper quirks). Am I right?

yeah exactly. zuko does implement several flows, but all within a common framework as much as possible. Compare that to, for example lucidrains, who does everything almost exactly like the paper (down to the name of the variables). I don't think that's necessary for this library.

[simonschnake]

Hi @francois-rozet,

I think such a library would be a fantastic idea and there are a lot of people (myself included) who would benefit from it. I would definitely recommend not to include the functionality in Zuko. Zuko has a lot of complexity on its own, using lazy transformations and so on. I don't see much benefit in including other classes of models and adding complexity.

I would think that most users would have specific research use cases where they would have to write custom models for their specific use case. Having primitives to build on would be a big help. Perhaps also including sample implementations that point to the dedicated papers. I would be more than willing to help write such a library.

Thanks for getting in touch and your work,
Simon

[francois-rozet]

Thank you @simonschnake. Definitely a separate library then. Tutorials to reproduce specific papers is a good idea!

[JuliaLinhart]

Hi everyone!

I love the idea! Zuko definitely makes the use of normalizing flows straightforward for many applications. Having the same kind of API for diffusion models would be amazing! I would definitely use it and would be happy to help! I especially like the idea of centralizing the different algorithms and features: i think there is a great potential of such a library becoming the standard tool for score based methods. I think we are getting to a point where there is a better understanding of such models and people want to start building on the baseline architectures, adding features, different samplers and schedulers, etc.

So yes, a new diffusion library would be great! For me integrating it to zuko doesn’t make any sense though.

[francois-rozet]

Thank you @JuliaLinhart. I think as well that the field has become quite mature, and a "standard" library would make it easier for newcomers to try things out.

[gerome-andry]

I see this a bit late, sorry!
Of course it is a very good idea! And I want to be part of it if you need any help!

I use to define over and over again my diffusion models in the same way, with sometimes small variations. I could definitely be a nice way to have everything packed in a clean library like zuko.

I agree with the others that a proper library, detached to zuko is a good idea to allow specific structure tailored to diffusion models as you suggest.

We can come back later to specific implementation-related topics, but I'm more of the team of having a unified and simple representation for diffusion models that could be specified to a task afterwards. So, about point 2, linking implementation to papers is obviously needed but under a simple unified implementation in my pov (to keep the user's experience straightforward).

[janfb]

Hi all,

@francois-rozet thanks for starting this discussion! I agree with what was said already. It's a great idea and a separate library would make sense.

At sbi we are now using zuko as our interface to flows and it is really nice to use! 👏
We also working on implementing score based SBI methods (sbi-dev/sbi#1015) and basing this on a dedicated library would be awesome. This would make it easier for users to access a larger variety of methods, and for researchers to try out new score-matching - SBI combinations.

[francois-rozet]

Thank you @janfb. Indeed, more variety is better for users and research!

[francois-rozet]

Thank you everyone for your answers. There is definitely interest for a diffusion models library 🤯 In the coming weeks, I will setup a repo in the probabilists organization and write a first draft for the interface. I invite you to open issues/discussions/PRs to discuss about anything (API, features, docs, ...).

Concerning the name of the project, I was thinking Azula would be fitting as she is the sister of Zuko 🔥 Tell me what you think!