more explaination and correct usage timestep_bias_strategy and other timestep bias argument

[zhuliyi0]

I have read the explaination in the source, but couldn't track it down where these arguments are used. Could you please provide more information of where and how they are used, and some recommendation/examples of what values can achieve what result?

[bghira]

there will be an article about it this weekend for huggingface. here are some basics.

This is how StabilityAI's training looks on a scatter plot, where we visualise timesteps per batch along the Y axis and steps over time via the X axis.

In other words, this is a solid block of blue because their scatter plot is fully covering all 1000 timesteps on every single batch size of 2048.

With a batch size of 128, we see major gaps in the timestep coverage on each batch. This is a visualisation of training on an A100-80G and those gaps are actually conservative, since the blue dots are fairly large. In other words, we have even less coverage over each training step than this.

Here is what it would look like if you were to visualise:
--timestep_bias_strategy=later --timestep_bias_portion=0.25 --timestep_bias_multiplier=5

The 'later' timesteps (higher numeric values) will be biased 5x higher than any other timestep in the schedule, while not entirely avoiding training of those timesteps.

If we increase --timestep_bias_multiplier from 5 to 20 then we'll see a chart more like this:

WHY?

The SDXL refiner was trained on just 200 of the 1000 total timesteps, rather than biasing them. The idea behind this is the same.

However, we are building potentially base models with this tool. Biasing the later timesteps has actually been pretty instructive for coherence, while the earlier timesteps have been apparently impacting fine-details:

Training an SDXL v-prediction variant on a 5x bias toward the earlier 25% of timesteps has led to this kind of weird monstrosity butterfly, but with some very good fine details - notably, there is some residual noise:

Taking that same checkpoint and then applying the inverse, a 25% bias toward the later 25% timesteps of the schedule resulted in this improvement:

It seems if the loss landscape of the later portions of the model are still crusty, that no matter how much you try and hammer on the early timesteps, at inference time, the model's validations will look like garbage.

But it's promising that we can potentially direct training toward the later timesteps and make more sweeping changes to the model, and direct training toward earlier timesteps for improving fine-details with very high quality training data.

This should ultimately result in less overfitting as the training run will not need to unnecessarily sample on unrelated timesteps.

[bghira]

sorry, my copy-paste went to hell there. i have fixed all the images.

here is another comparison, with earlier timesteps biased on the right, and later biased on the left

however, it's worth noting that i'm doing a major overhaul of the model, and that incoherence isn't the fault of training the early timesteps. it is how the model was already. just pointing out that training on later timesteps seems to improve larger features of the image.

[zhuliyi0]

looks like a very useful technique! I was guessing this must somehow related to the refiner. So you were saying the refiner was basically trained with the later 800 steps freezed, or in other words, infinitely biased towards the earlier 200 steps?

What about batch size? Does a small batch size still work?

[bghira]

the refiner seems to be "capable" of doing the early noise schedule but it doesn't do it very well. i don't know if that's because it "Never" saw it, or because it was fine-tuned on just the final inference steps.

small batch sizes have always been terrible but if you have a small dataset with a lot of visually-similar images, a large batch size might lead to overfitting.

[zhuliyi0]

can I change bias when continue from last checkpoint?

[bghira]

yes, it's got no impact in that way, unlike batch size or learning rate.

[zhuliyi0]

the argument "timestep_bias_begin" and "timestep_bias_end" requires --timestep_bias_strategy=range as the source says, but "--timestep_bias_strategy" says it only has three options ["earlier", "later", "none"]. Is "range" not supported yet? I added it to the choices, it seems to be running normally.

[bghira]

it's there, it's just missing in the help output

[zhuliyi0]

ok finially I found this old discussion.

I am in a training run, the model shows sign of converge. Validation image's composition is almost fixed, most things looks correct at first glance, saturation is increasing. Some small details needs to be fixed, and I am wondering if this is the right moment to use time step bias for earlier steps. If yes, should lr be decreased or just keep it as it is?

I would try and found out but would like to know more about the thing I am about to do. Thanks.

[bghira]

i think you are correct but it needs experimentation to know the best values