Do diffusion models take a long time to train?

Question

I am trying to train a diffusion model (from scratch in pytorch). UNet used is not anything too fancy, takes in images and time step as input for about 512 time steps. I am using learnable embeddings for 512 time steps. This works just fine for mnist but when moving to another dataset (my dataset is of motorbikes) it produces nothing meaningful even after 20k steps (batchsize 32). This is what i get after 20k steps on it.

Some experiments i tried somewhat contradict each other and i cant think of why this is happening.

Experiment:

1. overfit a small batch of my dataset. Result: model produces something intangible like the above picture
2. Have a dataset of about 10k images. Produces something like a bike if you really look for it. Trained this for ~28k steps (batchsize 32).
3. model works fine for mnist and cifar (even after running it for a short while).
4. Noticed there was noisy data which has nothing to do with bikes, so i manually got rid of them hoping my model would perform better. (reduced my dataset of 10k images to ~6.5k clean, quality images). Result: No changes, produces intangible images (just like first one) after training for ~55k steps (batchsize 32). This surprises me because i was expecting it to atleast be like the images produced from the noisy dataset.

I have not changed model hyperparameters when performing these experiments:

1. lr = 4e-4 with Adam
2. batch_size = 32 (for bikes)
3. Normalised images from [-1,1]
4. img_size = (3, 128, 128) # c, h, w
5. timesteps = 512
6. betas scheduler from 10e-4 to 0.02 (just like in the original paper)
7. using a cosine scheduler (this works a bit better than using a linear one)

I tried using self attention but that performs worse (not sure if its "worse" because it produces noise even after 12.5k steps on mnist (batchsize of 64).

I could not find a practical guide on training diffusion models, however from the information i gathered, the takeaway is that you have to train these forever and wont really work well if you run it for few hours on a mediocre gpu. I see that some of the sota diffusion models were trained on a few thousand days worth of gpu compute. Is my assumption correct ?

I am still not clear on why overfitting on a batch does not work here or why the clean dataset does not produce anything tangible. How do i know if my model is ever learning ? when do i know i should stop ?

Answer 1

Diffusion models are very data hungry. Without data augmentation that won't be enough to train it nor would I expect realistic images from a dataset that small for something so complicated. Diffusion models also are very bad at producing noisy images, so you are going to have to train it on a set of images that are less noisy.

My suggestion?

• More Data and make it clean. The more complicated, the more data you need. The original paper used millions to get realistic images. I would say at a minimum 20k is barely enough to see reasonable results without overfitting.
• Make sure you self attention is using convolutions instead of MLP's. This is what they call "attention" if you look at the implementation on the official DDIM repo.
• Adjust hyper parameters. Diffusion models can sometimes be very sensitive to hyper parameters. It's not easy like a VAE but not as hard as GAN. Try some more fiddling.

Answer 2

To assess whether your model is learning, you can monitor the loss function during training. If the loss is decreasing over time, it indicates that the model is making progress. Additionally, you can visually inspect the generated images at different training stages to see if they progressively resemble motorbikes more accurately.

You can employ techniques such as early stopping, where training is stopped if there is no improvement in validation loss over a certain number of epochs. Additionally, you can use validation metrics to assess the performance of the model on a separate validation set and stop training when the performance reaches a satisfactory level.