Add a bias_correction_v flag to scale_by_amsgrad to align with the original AMSGrad paper and Pytorch/tensorflow impl

[vvsvictor]

Resolves #1389

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[rdyro]

Wouldn't setting bias_correction_v=False not skip applying bias b2 entirely now?

Is the point to not use the bias at all or change the order of the operations as discussed here: pytorch/pytorch#142323

Pytorch applies this the bias in amsgrad to this day: https://github.com/pytorch/pytorch/blob/2164b661219ab0a76aa018e955ba3d8e8f99c083/torch/optim/adam.py#L509

But tensorflow does not (I think): https://github.com/keras-team/keras/blob/f6c4ac55692c132cd16211f4877fac6dbeead749/keras/src/optimizers/adam.py#L130-L150

[vroulet]

Thanks for the change!
Let's just quickly agree on the new argument name and we should merge.

[vroulet]

"match pytorch behavior" -> no, see the conversation.

So just say "set to False to match original AMSGrad paper"

[vroulet]

bias_correction_v is not a great name.

bias_correction_nu is already better

debias_nu may even be better but "bias_correction" as a boolean argument is already used in e.g. rmsprop (shame on me for that naming).

@rdyro what do you think?

[vroulet]

Wouldn't setting bias_correction_v=False not skip applying bias b2 entirely now?

Is the point to not use the bias at all or change the order of the operations as discussed here: pytorch/pytorch#142323

Pytorch applies this the bias in amsgrad to this day: https://github.com/pytorch/pytorch/blob/2164b661219ab0a76aa018e955ba3d8e8f99c083/torch/optim/adam.py#L509

But tensorflow does not (I think): https://github.com/keras-team/keras/blob/f6c4ac55692c132cd16211f4877fac6dbeead749/keras/src/optimizers/adam.py#L130-L150

I think optax original implementation is the one that makes most sense (doing the bias correction after taking the max does not make sense to me).
However one could also simply remove the bias correction see plots in #1389. It seems to potentially improve and most importantly it aligns with the paper.

[vroulet]

Looking at the paper, neither mu nor nu had bias corrections.
So to fully align with the paper give the option to remove both.
Namely have debias_mu: bool = True, debias_nu: bool = True for example. (or bias_correction_mu: bool = True, bias_correction_nu: bool = True)

Change the code accordingly