get_params returned from the optimizers feels weird.

[uduse]

Feels like get_params does nothing special but indexing an optimizer's state, and its behavior actually has nothing to do with how we initialize the optimizer. Is there a reason behind this? i.e., is there a special case where we should use the get_params returned from an optimizer's initialization instead of just index the optimizer's state?

[tomhennigan]

This is one of the design decision optax changes (params and opt_state are decoupled). It might be worth trying that library instead.

I'm not aware of any cases where params = get_params(opt_state) is preferable to having params separate from opt_state, however changing this in jax.experimental.optimizers would be API breaking so I would guess this will not happen.

[uduse]

Yeah, that makes sense. I would give optax a try!

[davisyoshida]

I saw this the other day and tried out optax at well, but found that it used more memory for my application. I have no idea why this is the case.

I'm using donate_argnums to to get the JIT to eat the params/opt_state arguments in both cases, but it seems to like the jax.experimental.optimizers.adam version more than the optax.scale_by_adam one for whatever reason.

Just an anecdotal issue I had, but figured it was worth sharing.

[mattjj]

Using get_params rather than indexing the state makes the optimizer API slightly more general. That is, thanks to get_params an optimizer need not store its state in any particular structure or with any particular order, or even be committed to just having one parameter value in its state; the state is totally opaque to the caller. For example, maybe the optimization algorithm's state contains not just one but many parameter vectors, and get_params selects the current best, or takes a convex combination of them. As another example, maybe the state is stored in some packed data structure which allows for more efficient execution than with a pytree container of arrays. (At one point we experimented with representing optimizer states as XLA runtime tuples rather pytree containers, meaning the cost to dispatch an optimizer step did not depend on the number of array buffers it represented. That kind of trick could be increasingly important in some large-scale regimes!)

That said, that generality might not actually be useful in a given context. That's partially why optimizers.py is just an example. If you don't like something about it, it's easy enough to fork and change. Or if you want a more robust, actively-developed library rather than an example, tools like optax are perfect!

[uduse]

@mattjj Thank you for your answer, that makes sense. 👍