What does the learning_rate parameter mean?

[nalzok]

For example, in the documentation for optax.adam, we have

learning_rate (Union[float, Callable[[Union[ndarray, float, int]], Union[ndarray, float, int]]]) – this is a fixed global scaling factor.

I noticed that learning_rate can be a function. How can a function be "a fixed global scaling factor"? What parameters should such a function take, and how does its return value affect the optimization process given that Adam has a learning rate schedule on its own?

[rosshemsley]

Hello @nalzok!

Thanks for your question. Unfortunately, this is a place where the generated documentation does not do a good job of reflecting the code.

If you take a look at https://github.com/deepmind/optax/blob/master/optax/_src/alias.py#L226

You will see that the actual type is ScalarOrSchedule.
If a scalar is provided, then the same value is used for every step.
If a Schedule function is provided, then we will store the number of steps taken so far within the optimizer state, and call the Schedule function with the current number of steps to get the learning rate to use for each step.

I agree the documentation could be clearer here. We can raise an issue for this.