Feat: Add decoupled, dimension-scaled weight decay option

[Koratahiu]

Background

Empirically, we can bypass the pitfalls of weight decay (WD) tuning by setting a single optimal value (WD * LR = 0.1).

We achieve this by scaling the weight decay by the width of the tensor, which scales linearly with the squared Frobenius norm.

In summary, the scaling rules are:

• Full Finetuning: WD / width
• LoRA: WD / width (Note: In theory, the LoRA-B factor has a higher Frobenius norm, so it should decay more heavily than the A-factor. But standard optimizers like Adam struggle to grow the B-factor so I recommend using spectral scaling with LoRA scaled WD).
• OFT: 2 * WD / block_size (The factor of 2 accounts for the skew-symmetric matrix having double the Frobenius norm, and block size is the width of constructed skew-symmetric matrix).

This approach simplifies WD tuning, ensures fair comparisons between layers and training methods, and is compatible with all optimizers.

New Option: scaled_wd

This option implements the scaling rules mentioned above and decouples the LR from the WD. When enabled, users can simply set WD = 0.1 and forget about it.

Additional Notes & Results

• Full Finetuning: My previous tests demonstrate that this rule works reliably. I applied it using a Centered WD = 0.1, which yielded the best SDXL finetune I currently use. Standard WD for full finetuning can be overly sensitive, as the norms were originally trained with very large batch sizes, often degrading quality in small-to-medium training runs.
• LoRA and OFT: Using a 0.1 scaled WD outperformed my hand-tuned WD values on the very first attempt.

[khanghy1000]

Edit: LR = 1e-4 is too low for spectral normalization.

I used scaled_wd to train a lora for anima in sd-scripts and it barely learned anything. It seems like the weight decay was too aggressive. Is wd = 0.1 the correct value for scaled_wd or should i use a lower wd?

Training settings:

• Optimizer: Adamw_Adv
• Initial LR: 1e-4
• LR scheduler: REX (1 cycle)
• Global batch size: 2
• weight_decay: 0.1
• cautious_wd: True
• betas: (0.9, 0.99)
• compiled_optimizer: True
• use_atan2: True
• spectral_normalization: True
• scaled_wd: True

[Koratahiu]

I used scaled_wd to train a lora for anima in sd-scripts and it barely learned anything. It seems like the weight decay was too aggressive. Is wd = 0.1 the correct value for scaled_wd or should i use a lower wd?

Training settings:

• Optimizer: Adamw_Adv
• Initial LR: 1e-4
• LR scheduler: REX (1 cycle)
• Global batch size: 2
• weight_decay: 0.1
• cautious_wd: True
• betas: (0.9, 0.99)
• compiled_optimizer: True
• use_atan2: True
• spectral_normalization: True
• scaled_wd: True

What is the LoRA rank so I can test it?
Also, spectral normalization requires tagging the parameters so it can distinguish between the LoRA A and LoRA B factors.
Something like:

for param_name, p in lora_module.named_parameters():
    if param_name.endswith("lora_a.weight"):
        p._is_lora_A = True

[khanghy1000]

What is the LoRA rank so I can test it?

It's a rank 8 LoRA.

Sorry, I think it was because the LR was too low for spectral normalization. I increased it to 1e-3 and the LoRA seems to be learning now. The learning is still weaker than it is without scaled wd, but it does learn.