The College Dropout x The New Norm(al)

Hey, wanted to talk a little about a couple of mystical kohya_ss parameters, and how doing the math to set them could reduce the subjective guesswork in training a Lora. When things are a little close, or the Loras from two epochs might compete with each other based off the prompt of checkpoint, testing those results can take a lot of time.

One tool I use heavily is Tensorboard, which was a game changer as looking at loss/epoch helps identify candidates quickly. It's a time-saver to lower your candidates down to a handful.

I also started using the Prodigy optimizer almost exclusively, as it adjusts the learning rates on the fly and makes training faster and, again, more objective. One of the optimizer arguments is called d_coef, which controls Prodigy's adaptive learning rate (d) and can dampen (< 1) or speed up learning (> 1). If it trains too slowly, you will train it on things you don't care about (like backgrounds); if you train it too fast, it will overtrain. I have read the recommended values are between 0.5 and 2 but not much more than that.

I guess curiosity got the best of me one day and I did some Googling and found this Reddit thread. The interesting comment starts here: "They are looking at the tensorboard to draw conclusions about the training. And based on that change the "d_coef" for prodigy."

So I clicked though the Github link and skimmed the PR about dropout and max norm regularization, then read the aforementioned comment: https://github.com/kohya-ss/sd-scripts/pull/545#issuecomment-1831901651

"I have trained several Lora in the last weeks and I have noticed that a stable, good quality and yet flexible Lora is made with Prodigy, cosine scheduler and if max_norm/keys_scaled does not go above 2 (sometimes even 2 is strong) when training a person. For styles I usually leave it stronger, up to 6. If the max_norm/max_key_norm doesn't reach 1, it will be under-trained. For Prodigy, this is best done by changing the d_coef value."

(In the Bmaltais Kohya GUI, these features are called scale weight norms and network dropout, and there is a short description of each in the GUI.)

I did a few experiments one weekend trying to see if I could linearly adjust d_coef to time max_norm/keys_scaled to hit 2 right at the lowest error-rate epoch. They were largely a failure, because changing d_coef didn't predictably affect max_norm/keys_scaled. So while I absorbed the knowledge for future training, I forgot about trying to get it perfect.

This weekend, I accidentally set scale weight norms to 5, then wasted a day trying to figure out why all my runs were over-trained. Once I figured it out, I went back to those articles, and noticed something in the very first post in the thread:

"- Max Norm suggested setting = 1 (You can also set it high enough to never trigger ie 10 to watch Tensor Board and see where a good point to set it at might be)"

What does that mean exactly?

I looked at the failed runs, where the keys are never scaled. And I remembered that when the keys start to be scaled at 1, when scale weight norms is set to 1, that's the indicator the Lora is almost trained. This commenter seems to think you should ignore the loss/epoch , because the max_norm/max_key_norm being ~2 is a better indicator. OK, but I thought, what if you could get both? A Lora that's had just enough training, and has low loss/epoch for prompt adherence?

Here is what I tried:

1. I set my parameters like usual. Because I noticed previously that if you don't mess with max steps/epochs, batches, dim, or alpha, the loss/epoch graph will follow the exact same path. So I decided to train an XL character Lora using dim = 64, alpha = 32, and d_coef = 1. I set network dropout to 0.3.

2. You need to understand batch sizes and buckets and epochs. If you're not sure, you can start the training and kill it after it tells you your batches per epoch. I had a batch size of 4 and batches/epoch of 17. I was using regularization images, so I decided to aim for 2500 steps. I wanted some cushion for a couple of extra epochs, so I settled on 40 epochs and 680 max steps (actually 2,720, when accounting for the batch size of 4).

3. I set scale weight norms to 10 so it will never scale.

4. I let it run and did something else for a bit. When it was done, I checked the loss/epoch Tensorgraph like normal and while there was a good dip in the mid-20s, I aimed for epic 32, which would be at 544 steps.

5. I looked at the max_norm/keys_scaledgraph, where the slope of the plot was almost linear. I started by looking for the value on the y-axis where the x-axis (steps) is at 544. Because I don't want to hit epoch 32 right when the keys are scaled, because I can squeeze out some additional training, I backed up a couple of epochs. That's 34 steps in this example, to 510 steps. The y-value at 510 is 1.346, so I set scale weight norms to that value. What this means is that on the next run, instead of climbing at a roughly steady state, max_norm/keys_scaled will level off at 1.346 and then get two epochs of extra training in, as the max_norm/keys_scaled will escalate pretty quickly.

6. Luckliy, scale weight norms does work pretty linearly. If scale weight norms goes too high or not high enough, I can pick the spot where it's around 2, calculate the delta in steps from where we want to end up (544 steps), then repeat step 5. So if the Lora is juuust over-trained by 15 steps, I can look at the y-value of the test run again, and change scale weight norms to whatever y was when x was 495 (hypothetically something like 1.32).

I did this with three Loras total, and they came out pretty good. One of them I had just trained the day before, and the new Lora was as good or better. It was very similar at reproducing the character but had better adherence to the prompt and looked more like the image without the Lora applied. It seemed to completely eliminate the guesswork, even if it took more time to do 2 or 3 trainings.

I'll try to update this with photos, but let me know what you think.