Implementation of 'Score-based Diffusion Models in Function Space'
This repository implements the method, denoising diffusion operator (DDO), proposed in Score-based Diffusion Models in Function Space, i.e., a function-space version of diffusion probabilistic models, using JAX and Flax.
Important
The implementation does not strictly follow the experimental setup in the paper (since the paper itself uses a different one for each experiment). Specifically, the U-net neural operator (U-NO) as well as the sampling are customized and simplified. Our U-NO implementation just uses spectral convolutions for up- and down-sampling of input dimensions. We use the VP-parameterization of DDPM; hence we don't use the score-matching loss in NCSN but a conventional SSE loss. We consequently also don't use Langevin dynamics for sampling, but the sampling proposed in DDIM.
If you find bugs, please open an issue and report them.
The experiments folder contains a use case on MNIST-SDF. For training on 32x32-dimensional images from the MNIST-SDF dataset, call:
cd experiments/mnist_sdf
python main.py \
--config=config.py \
--mode=train \
--model=<uno|unet> \
--dataset=mnist_sdf \
--workdir=<dir>Then, sample images via:
cd experiments/mnist_sdf
python main.py \
--config=config.py \
--mode=sample \
--model=<uno|unet> \
--dataset=mnist_sdf \
--workdir=<dir>Below are DDIM-sampled images from the DDO when either a UNet or a UNO is used as score model (a DDO with a UNet is just a DDPM). The UNet parameterization yields high-quality results already after 20 epochs or so. The UNO works worse than the UNet when 32x32-dimensional images are sampled and takes significantly longer to train. When sampling 64x64-dimensional images it mainly produces noise.
To install the latest GitHub , just call the following on the command line:
pip install git+https://github.com/dirmeier/ddo@<TAG>Simon Dirmeier sfyrbnd @ pm me