Inconsistent description of AMSGrad with code

[Tony-Y]

📚 The doc issue

pytorch/torch/optim/adam.py

Lines 469 to 476 in 0bd7b7a

	if amsgrad:
	# Maintains the maximum of all 2nd moment running avg. till now
	torch.maximum(max_exp_avg_sqs[i], exp_avg_sq, out=max_exp_avg_sqs[i])
	# Use the max. for normalizing running avg. of gradient
	denom = (max_exp_avg_sqs[i].sqrt() / bias_correction2_sqrt).add_(eps)
	else:
	denom = (exp_avg_sq.sqrt() / bias_correction2_sqrt).add_(eps)

In the code, the bias correction term $1-\beta_2^t$ is used after the max operation. However, the documentation describes that it is used before the max operation:

Suggest a potential alternative/fix

cc @svekars @brycebortree @sekyondaMeta @AlannaBurke @vincentqb @jbschlosser @albanD @janeyx99 @crcrpar

[Tony-Y]

$\widehat{v_{t}}$ may be used instead of $\widehat{v_{t}}^{max\ \ }$ in order to simplify the description:

[Tony-Y]

By the way, optax implements the AMSGrad algorithm described in PyTorch 2.5 docs:
https://github.com/google-deepmind/optax/blob/78fe20e75c79b2f615bcd9e8a653af9d988f8992/optax/_src/transform.py#L356-L361

[janeyx99]

To be clear, the actual math is equivalent between the docs and the code, right? The v_max thing is consistent in docs + computation even though the order is different, since dividing by a constant doesn't affect the max, yes?

By the way, we don't actually have guarantees on the code and the docs being exactly matching (as I'm sure you'll find other instances of this) and it is okay for them to not match to allow the code to take performance wins and for the docs to be more readable. In this case, your suggestions are more readable so let me go review ur PR 😛

[Tony-Y]

Not equivalent. The bias correction term is not a constant. It depends on t.

>>> t = [1, 2, 3]
>>> v = [5, 4, 3]
>>> b2 = 0.5
>>> vmax = 0
>>> for i in range(3):
...     vmax = max(vmax, v[i])
...     
>>> vmax / (1-b2**t[i])
5.714285714285714
>>> vmax = 0
>>> for i in range(3):
...     vmax = max(vmax, v[i] / (1-b2**t[i]))
...     
>>> vmax
10.0

Edit
I fixed a bug of the former calculation.

[janeyx99]

ah........in that case we should make sure our calculation is aligned with the AMSGRAD paper if possible.

[Tony-Y]

The PyTorch implementation is similar to Tensorflow one:
https://github.com/keras-team/keras/blob/f6c4ac55692c132cd16211f4877fac6dbeead749/keras/src/optimizers/adam.py#L130-L150

The AMSGrad paper does not mention any bias correction.