Adam is slower than sgd+momentum on flax

[lucasliunju]

Hi I find Adam is slower than sgd+momentum, which increase 20% time cost compares with sgd+momentum. I'm not sure whether that is normal. Could you give me some advice on that?

Thank you!

[matthias-wright]

Hi @lucasliunju! Some clarifying questions:

• Are you talking about the time difference for one optimization step or the difference in time for your optimization process to converge?
• Are you using flax.optim or optax? Because as far as I understand optax is now the redommended optimizer API (see #1095).

[lucasliunju]

Thanks for your reply.

1. The training time of each step and not the convergence time.
2. I am using flax.optim.

I think the main difference between Adam and SGD+Momentum is that Adam store more variable and use more memory. I'm not sure whether that will cause so big time cost.

Thank you!

[matthias-wright]

Yeah, Adam is essentially SGD + Momentum + RMSprop. So in addition to storing the exponential moving average of the gradients for momentum, the exponential moving average of the squared gradients is also stored for RMSprop.
At each update step, the exponential moving average of the gradients is divided by the square root of the exponential moving average of the squared gradients (ignoring bias correction), so there are some more computations. Hence it is not surprising that the Adam step takes more time.
I just tested it with a small example and I am observing around 10% time increase with Adam in comparison to SGD+Momentum.
Maybe you can also create a small reproducible example from your code.

[jheek]

Please keep in mind to use jit and .block_until_ready for benchmarks. see link