LoRA weight scaling

[dxqb]

Very interesting paper, and thank you for publishing code with it!

I have a question related to LoRA weight scaling:
The original LoRA paper, https://arxiv.org/pdf/2106.09685 proposed scaling the LoRA weights by alpha / rank.
This is still used in modern code, as it has been shown that this roughly allows you to change rank without changing learning rate. For example diffusers: https://github.com/huggingface/diffusers/blob/3c8b67b3711b668a6e7867e08b54280e51454eb5/src/diffusers/models/lora.py#L230
OneTrainer, diffusion model trainer: https://github.com/Nerogar/OneTrainer/blob/6276c512c2ff6ad50e74eb4617274ed2f44bdcee/modules/module/LoRAModule.py#L323

I wonder how this applies to your proposal to vary the effective rank by timestep.

I haven't found any mention in your paper except the use of a constant learning rate, but in your code you seem to have removed this scaling :

T-LoRA/tlora/model/lora.py

Line 33 in ba41983

def forward(self, hidden_states, mask=None):

Could you address this?
Doesn't this mean that you effectively apply larger parameter updates to the higher timesteps, than you would have done if weight scaling was used, because the effective rank after masking is lower?
Parameter update size (effective learning rate) has a very high impact on overfitting in finetuning diffusion models.

[V-Soboleva]

Dear @dxqbYD,

Thank you for your interest in our paper and for your question!

We did not use scaling because we based our implementation and experiments on the LoRA training examples for Stable Diffusion XL provided in the Diffusers library. In these examples, the lora_alpha parameter is set equal to the rank by default (this is also the case for LoRA training examples with other models). As a result, the scaling factor, defined as lora_alpha / rank, is always equal to 1.

You can find the relevant code here:
https://github.com/huggingface/diffusers/blob/3c8b67b3711b668a6e7867e08b54280e51454eb5/examples/dreambooth/train_dreambooth_lora_sdxl.py#L1208

Additionally, while a smaller learning rate could potentially help reduce overfitting, our experiments primarily focused on comparing the models under consistent conditions. As such, we did not manipulate the learning rate and instead used its default value.

If we understand your proposal correctly, you suggest complementing rank masking with an adaptive adjustment of the learning rate (by modifying the scaling factor). Specifically, this would involve using a higher learning rate for noisier timesteps and a lower learning rate for less noisy ones. While this approach is intriguing, it somewhat diverges from the motivation of our paper, where we advocate for stronger regularization at earlier timesteps compared to later ones.

However, the idea that different timesteps and ranks require different learning rates is interesting. We are leaving it for further research, as the patterns for diffusion models may differ from what was proposed in the original LoRA designed for LLMs.

[dxqb]

We did not use scaling because we based our implementation and experiments on the LoRA training examples for Stable Diffusion XL provided in the Diffusers library. In these examples, the lora_alpha parameter is set equal to the rank by default (this is also the case for LoRA training examples with other models). As a result, the scaling factor, defined as lora_alpha / rank, is always equal to 1.

A scale factor of 1.0 can work of course, if you do your hyperparameter search for this specific rank. However, it's a known fact among practitioners that as soon as you change LoRA rank, you have to significantly change your learning rate - or keep your LR the same but approximate this change by keeping your LoRA alpha unchanged while changing your rank - therefore changing your LoRA scale factor.

If we understand your proposal correctly, you suggest complementing rank masking with an adaptive adjustment of the learning rate (by modifying the scaling factor). Specifically, this would involve using a higher learning rate for noisier timesteps and a lower learning rate for less noisy ones. While this approach is intriguing, it somewhat diverges from the motivation of our paper, where we advocate for stronger regularization at earlier timesteps compared to later ones.

However, the idea that different timesteps and ranks require different learning rates is interesting. We are leaving it for further research, as the patterns for diffusion models may differ from what was proposed in the original LoRA designed for LLMs.

I think it is important to include this in your research:
Worst-case, your experimental results could be hyperparameter bias, because you change a hyperparameter (rank) that is highly dependent on another hyperparameter (learning rate).
Best-case, and somewhat more intuitive, using a variable scale factor might even significantly improve the benefits of T-LoRA and contradict your finding that r_min at 50% / r=32 has the best result:
This finding is somewhat surprising, because practitioners often use ranks below 32 with good results. So, if LoRAs with ranks below 32 can work well, why would T-LoRA, with the purpose to better adapt effective rank, not work well below 32 while keeping a higher rank for those timesteps that benefit from it? It should be the opposite, that you can use a lower r_min with T-LoRA than you can use with regular LoRA.

So my intuition would be that - because you apply no scaling and a fixed LR - the effective learning rate (parameter updates) for ranks below 32 becomes ineffective. Anyone who changes rank without T-LoRA from 64 to 32 or even lower would change their learning rate (or keep alpha unchanged). If you applied scaling based on effective rank, r_min below 50% might be possible and show the strength of T-LoRA.

[dxqb]

edited with some corrections / clarifications