Dreambooth broken, possibly because of ADAM optimizer, possibly more.

[affableroots]

I think Huggingface's Dreambooth is the only popular SD implementation that also uses Prior Preservation Loss, so I've been motivated to get it working, but the results have been terrible, and the entire model degrades, regardless of: # timesteps, learning rate, PPL turned on/off, # instance samples, # class regularization samples, etc. I've read the paper, and found that they actually unfreeze everything including the text embedder (and VAE? I'm not sure so I leave it frozen), so I implemented textual inversion within the dreambooth example (new token, unfreeze a single row of the embedder), which improves results considerably, but the whole model still degrades no matter what.

Someone smarter than me can confirm, but I think the culprit is ADAM:

diffusers/examples/dreambooth/train_dreambooth.py

Line 426 in f1484b8

optimizer_class = torch.optim.AdamW

My hypothesis is that since ADAM tries to drag all weights of unet etc. to 0, it ruins parts of the model that aren't concurrently being trained during the finetuning.

I've tested with weight_decay set to 0, and results seem considerably better, but I think the entire model is still degrading. I'm trying SGD next, so, fingers crossed, but there may still be some dragon lurking in the depths even despite removing ADAM.

A point of reference on this journey is the JoePenna "Dreambooth" library which doesn't implement PPL, and yet preserves priors much much better than this example, not to mention it learns the instance better, and is far more editable, and preserves out-of-class rather well. I expect more from this huggingface dreambooth example, and I'm trying to find why it's not delivering.

Any thoughts or guidance?

EDIT1A: SGD didn't learn the instance at 1000 steps + lr=5e-5, but it definitely preserved the priors way better (upon visual inspection. The loss really doesn't decrease much in any of my inversion/dreambooth experiments).

EDIT1B: Another test failed to learn using SGD at 1500 steps + lr=1e-3 + momentum=0.9. It might be trying to learn, but, not much. Priors were nicely preserved though still.

EDIT1C: 1500 * lr=5e2 learned great, was editable, didn't destroy other priors!!!

EDIT2: JoePenna seems to use AdamW, so I'm not sure what's up anymore, but I'm still getting quite poor results training with this library's (huggingface's) DB example.

[osanseviero]

cc @patil-suraj @apolinario

[patil-suraj]

Hi @affableroots , thanks for the issue!

What do you mean by the whole model still degrades ?

The dreambooth paper uses imagen model, which keeps the text encoder frozen. The VAE should also be frozen.

As far as I can see this https://github.com/JoePenna/Dreambooth-Stable-Diffusion also uses prior preservation. Did you try training with this and did it work ? I would be happy to investigate.

Also note that, for dreambooth we need to play around a bit with hyperparamters i.e the number of steps, learning rate, number of prior images etc. Could you please post more details

• How many images are you using for training
• How many priro images
• The number of training steps
• The instance prompt and class prompt
• And if possible the training images that you are using.

This will help us debug the issue better. Thank you.

[affableroots]

If I understand correctly, this paper says they finetune the entire model, although in my experiments I agree and have kept the VAE frozen. I guess the following could be saying they only finetune all layers down stream of the text embeddings?

---

RE JoePenna, I think his repo is an example of Textual Inversion + unfrozen unet. Quoted from his readme:

We're [now] realizing the "regularization class" bits of this code have no effect, and that there is little to no prior preservation loss.

So, out of respect to both the MIT team and the Google researchers, I'm renaming this fork to: "Unfrozen Model Textual Inversion for Stable Diffusion".

For an alternate implementation that attempts prior loss preservation, please see "Alternate Option" below.

This repo currently gives far better results than the huggingface version, but I think the HF version has more promise once fixed.

---

RE my params, I've tried ~50ish permutations within the following bounds, but in general I've tried to stick close to the example params in your readme, and the params in the DB paper.

• Batches: 1-2
• 8bit optimizer on/off
• Number of images: 5 - 50
• Number of prior images: 200-1500
• Number of steps: 500 - 5000
• Learning rate: 1e-7 to 1e-4
• Prior Preservation Loss on/off
• Prompts are always of the form: "a photo of CLASS" and "a photo of TOKEN CLASS", where TOKEN is sks (struggles to train away from rifles), or random three letter words, or, more recently, I added textual inversion to your example to add a token to the tokenizer.
• I've tried many types of images, but have settled on a test case of simply learning faces that aren't already in SD, eg, my face.

---

RE degradation, I mean the whole model slowly becomes unusable over the course of learning. For instance given the stock params of 500 steps * lr=5e-6, it poorly learns the instance face, and all the class images visually degrade somewhat. It forgets things it knows already, and generally turns to visual noise, at any learning rate, and for any number of steps, in or out of the finetuning set. Very low learning rates, eg 1e-7, slow the process, but then you're not learning your instance either.

For inference I always test with: "a photo of token class", " a black and white sketch of token class", "a photo of johnny depp". So I've got in there: the exact prompt I trained with, a prompt that tests generalization, and a non-instance prompt to see if it's forgetting other faces. Johnny Depp degrades over the course of learning.

---

Given my limited understanding, could this all be due to weight regularization performed by Adam, decaying the weights? I seem to have better results using SGD, although it learns slowly.

[jslegers]

I'm still getting quite poor results training with this library's (huggingface's) DB example.

I've been using ShivamShrirao's fork of this repo for several days now. With my own Google Colab notebook (which adds additional configuration options & streamlines the whole process), I've managed to generate some pretty decent renders of my younger self...

Some of the renders generated

Actual photos of me used as input for the training process

[affableroots]

@jslegers Thanks for the follow up! I've had results look fairly similar, but I think everythings still degrading at least on my end. For instance:

• "a photo of" type prompts look a bit overcooked for the trained instance, as well as out-of class instances (unless my lr is low enough to not learn the instance)
• How do results look for random actors, compared to pre-training? Mine tend to degrade considerably, and often look "overcooked", ie that look things get at real high CFG.
• And how about using your likeness in non-portrait situations like, playing a sport, or an instrument, etc.?

[affableroots]

Oh, another symptom that I think indicates something is wrong, is I have not been able to overtrain. What I mean is, given a very high lr and/or high step count, everything sort of degrades without looking like my new instance. Faces look melted, and drift away from their priors, eg actors faces are no longer their faces, and they probably "melt" and become very indistinct, as faces. I notice photos degrade quicker than artwork.

Contrarily, I feel like JoePenna's version is easy to overtrain. Sure the trained concept leaks into everything, but I'm using this as a sort of diagnostic tool. In his version, an overtrained concept looks unmistakably like that concept (even if it increasingly fails to generalize)

Is there no merit to the weight decay hypothesis? That still sticks out in my mind. The way I see it is if you have a whole universe of latent space, but you only update the small corner that looks like your face, but regardless you regularize the entire thing towards 0, you're going to degrade everything, which would explain why overtraining doesn't work, generalization, learning a concept in the first place, etc.

Given my hypothesis, with a low step count you could learn your concept without degrading things too much, as @jslegers and others (and myself) have demonstrated, but it still doesn't seem like a quality finetuning for all the reasons I've mentioned, IE how the model degrades.

Or I could be crazy and running this all wrong?

[jslegers]

but I think everythings still degrading at least on my end.

For my use case, I noticed that the diversity with the class "man" started declining significiantly. At some point, all random class images generated during the training process look identical.

I haven't really looked much into the degrading of my "environment", though, as I was too focused to trying to fix my likeness. But it does indeed appear to be degrading as training increases.

How do results look for random actors, compared to pre-training? Mine tend to degrade considerably, and often look "overcooked", ie that look things get at real high CFG.

Is this that "overcooked" look you're talking about?

Nightmare fuel, if you ask me...

And how about using your likeness in non-portrait situations like, playing a sport, or an instrument, etc.?

I haven't really invested much time in doing full body renders, but my best models did produce some OK results :

he way I see it is if you have a whole universe of latent space, but you only update the small corner that looks like your face, but regardless you regularize the entire thing towards 0

That seems to be what's happening, yes.

As a sidenote, I'm struggling to find a way to build upon previous trainings. When I have a general "johnslegers" concept, for example, I'd like to build upon that training be training eg. the concept "close-up or johnslegers" or "detailed portrait of johnslegers", as to refine the model further. But this doesn't seem possible. Attempts to achieve this by tinkering with my class prompt, instance prompt, class name or concept name generally leads to a loss of previous training achieved and serious degrading the quality of my models.

[affableroots]

Thank you for running that again! It seems we have similar results, though, your generalization attempts do look better than what I've been achieving, even given the stock params from the readme. And yes, that's the overcooking I'm talking about. That seems specifically like certain latent values are... clipping, if that's the right term.

try to refine the model further

Ya imagine that's impossible if the entire latent space is getting regularized and forgetting everything it learned in its original training. I've turned off weight_decay and while it seems to help, the whole latent space gets degraded anyway.

So then I thought maybe the Prior Preservation Loss was affecting things incorrectly, and I turned that off. So, ostensibly it should become a run of the mill finetuning of just unet, I think. I would expect that training without PPL would simply, for instance, pull all faces toward your face, but not destroy everything. But, all faces somehow just get cooked!

Hypothesis: I'm strongly suspecting weight decay is an issue and some other hidden factor.

[jslegers]

Ya imagine that's impossible if the entire latent space is getting regularized and forgetting everything it learned in its original training.

Probably.

I guess I better just start all over with the original SD model and play around with different image files & config settings to reproduce at minimum the quality I achieved to far than continue tinkering with existing models...

[patil-suraj]

@affableroots In the paper by all layers they mean all layers of the unet model trained, the rest i.e the text encode is kept frozen.

Looking at the discussion and also from my experiments the model does run into overfitting and catastrophic forgetting. To me this seems likely as we finetune the whole model with few images. For now I don't know any way to avoid this, I am going to discuss this with the author and let you know if we find anything.

To sum it up, I'm not sure if the script is broken, it's rather that the model quickly overfits.

[affableroots]

I appreciate your attention on this @patil-suraj!

[jslegers]

Looking at the discussion and also from my experiments the model does run into overfitting and catastrophic forgetting. To me this seems likely as we finetune the whole model with few images. For now I don't know any way to avoid this, I am going to discuss this with the author and let you know if we find anything.

I suspect overfitting issues and other related issues can be reduced by finetuning your number of training steps, class prompt, number of variety of class images generated, quality and variety of your input images, etc.

I've been experimenting with various strategies : use just a handful of pics, use many pics, auto-generate the class pics, choose class pics myself, do only 250 training steps per run, do 2000 training steps each time, etc.

My results after just a single run from the original SD model vary significantly depending on various of these criteria and I have yet to determine any pattern.

We should develop some best practices on how to optimize the model we're creating while keeping degradation of the latent space minimal. If anyone should be able to provide some info we all seem to be missing, that's going to be the original author...

[affableroots]

While we're exploring the tuning of DB from all angles, I'll mention some recently decent results. I'm using my modified Dreambooth + Textual Inversion:

• 5 new tokens+embeddings
• 2000 steps
• lr5e-5 for text embeddings
• lr5e-6 for unet
• weight decay=0.0
• prior loss weight=0.1
• no. class images=1500

I will say that faces are still getting cooked, so, I'm going to play with this some more.

My theory is that if you combine DB + TI, you can ask TI to learn more strongly without risking damage to the rest of the model (because you're only learning brand new embeddings that don't cause other domains to forget), but still train unet to help out a little, where TI is struggling.