What means the step in the update function?

[jingtianweidizzf]

Dear all,

I also have a question about the jax.experimental.optimizers module.

In the update_fun(step, grads, opt_state) there is an input variable called step and in the example given, the step just means the number of the steps.

But if I'm training a 2-layer vanilla CNN on MNIST for num_epoch epochs and in each epoch, there are num_batch batches. So when calling that update function update_fun, which one shall I use as the input variable step? The epoch number of the whole training process, or the batch number in one epoch, or the batch number totally (epoch_now * batch_num + batch_now ).

And will this step input make a great difference to my training process?

Thank you in advance for your time! :)

Best,
Zifan

[hawkinsp]

I'd suggest reading the code of the optimizers module: it's pretty readable on the whole.

Step numbers are used for at least two reasons:

1. for learning rate schedules, where parameters like the learning rate are varied as training progresses. For example, if you look at the implementation of sgd (https://cs.opensource.google/jax/jax/+/master:jax/experimental/optimizers.py;drc=b260468b51efff40183796e98b844314b66f7686;l=246 ) it computes a step_size based on the current step number.
2. some optimizers use the step number explicitly. For example, the ADAM optimizer uses it in its bias correction term (https://cs.opensource.google/jax/jax/+/master:jax/experimental/optimizers.py;drc=b260468b51efff40183796e98b844314b66f7686;l=409).

These are both standard and not specific to JAX; the only slightly unusual thing about the JAX optimizers is that they make this step number calculation explicit and under your control.

I'll also note that Optax has largely subsumed jax.experimental.optimizers for many of our users (https://github.com/deepmind/optax ), so you may wish to look at that library also. Hope that helps!

[jingtianweidizzf]

Hi Peter,

thanks a lot for your reply! That really helps!

Best,
Zifan