Why not 1-sqrt for WSD decay?

[slashedstar]

In the section 4 of the paper referenced for Warmup-Stable-Decay learning rate scheduler (https://arxiv.org/abs/2405.18392) they mention 1-sqrt being the best cooldown function, on page 20 they further demonstrate how it outperforms other functions so why the choice of cosine for WSD cooldown? https://github.com/kozistr/pytorch_optimizer/blob/main/pytorch_optimizer/lr_scheduler/wsd.py#L33

[kozistr]

@slashedstar hi. thanks for raising this. At that time, as far as I remembered, I first looked into the MiniCPM experiment note and the PR in Transformer, and I just added it here right away without checking the WSD paper honestly.

Of course, We could add 1-sqrt and other cooldown functions too, and make 1-sqrt cooldown function as a default one cause it's a best choice followed by the paper.

I'll work on it sometime later and thanks again for asking :)

[kozistr]

@slashedstar hi. I've just merged this feature (and also more cooldown types) into the main branch, and it will be included in v3.6.1. thanks!