[Scheduler] Add Variational Diffusion Models (VDM) scheduler

[hummat]

What does this PR do?

This PR adds a discrete and continuous time scheduler based on the Variational Diffusion Models (VDM) formulation, i.e. expressing the diffusion process via the signal-to-noise ratio.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

@yiyixuxu

[yiyixuxu]

thanks for the PR!
can we see some examples with this scheduler?

[yiyixuxu]

also can you provide more information? when would you recommend this scheduler and how does it compare to the default ones?

[hummat]

I implemented it primarily to reproduce results from the original Variational Diffusion Models paper and the Locally Attentional SDF Diffusion for Controllable 3D Shape Generation paper. I didn't use it to train on image data, so I can't provide any generations.

Performance wise it sits somewhere between the EDM and DDPM formulation I would say. It outperforms other schedulers in terms of log-likelihood.

The continuous time formulation doesn't require to chose the number of train timesteps, so one hyperparameter less. It could also serve as an example of how to incorporate continuous time into the other schedulers.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[vladmandic]

what is the status of this pr?

[hummat]

@yiyixuxu

[yiyixuxu]

@hummat woud you be able to provide a comparison? like for sdxl, output with vdm, ddpm and dpm with same number of steps etc
otherwise I can look into this when I have a bit more time!

[vladmandic]

quick grid compare with few other samplers
imo, has a lot of potential because its so different, but right now math is a bit broken.

first, need to add order: int = 1 to class init as its looked up by pipelines even if its not used by sampler

but...there is something weird about its sample range - at default settings, its clipping so badly, its not even resolving.
the grid above is with clip_sample_range=2.0 and its not perfect.

looking at sample min/max for each timestep (before clamp), its massive!

{'timestep': 0.800000011920929, 'min': -17.625, 'max': 22.375}
{'timestep': 0.7330000400543213, 'min': -12.4375, 'max': 18.25}
{'timestep': 0.6670000553131104, 'min': -14.625, 'max': 16.125}
{'timestep': 0.6000000238418579, 'min': -15.25, 'max': 19.125}
{'timestep': 0.5330000519752502, 'min': -10.125, 'max': 14.5}
{'timestep': 0.46700000762939453, 'min': -9.25, 'max': 12.8125}
{'timestep': 0.4000000059604645, 'min': -7.78125, 'max': 8.125}
{'timestep': 0.3330000042915344, 'min': -5.8125, 'max': 7.40625}
{'timestep': 0.2670000195503235, 'min': -5.625, 'max': 4.1875}
{'timestep': 0.20000000298023224, 'min': -3.453125, 'max': 4.25}
{'timestep': 0.13300000131130219, 'min': -3.046875, 'max': 3.296875}
{'timestep': 0.06700000166893005, 'min': -2.671875, 'max': 2.921875}
{'timestep': 0.0, 'min': -2.046875, 'max': 2.046875}

i'm not an expert here, but to me it looks like its all about initial step as range does decrease with steps, it just starts too high.
just to check, i did the same for DDIM and initial range is about 1/3rd

[hummat]

Thanks for the great experiment @vladmandic! To be honest I'm surprised that it works at all as the current implementation follows the VDM paper where models are expected to be trained on the noise level (or equivalently the (log) signal-to-noise ratio) similar to the models from the EDM paper instead of directly on the time steps.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.