[advoptm] Decoupled, Dimension-Scaled Weight Decay

[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 Weight Decay

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.

Includes: #1344

[dxqb]

is this the PR you wanted to add no-weight-decay-for-bias? or is this one ready?

[Koratahiu]

is this the PR you wanted to add no-weight-decay-for-bias? or is this one ready?

Yeah, I already added it.
As for being ready, I think so. The final step is to publish the stable version, which I will do.
I tested it for OFT and LoRA, and the 0.1 scaled weight decay worked out of the box.

[dxqb]

put back to "Ready" then please