Implementation of MoMo algorithm

[fabian-sp]

Upon suggestion by @fabianp I implemented the MoMo algorithm. MoMo is esentially a Polyak step size for SGD with momentum and for Adam (see https://arxiv.org/abs/2305.07583).

The Rosenbrock and least squares tests are passing locally.

I have still a few questions as this is the first time I am implementing in Optax:

• MoMo needs in each iteration the latest batch loss passed into update_fn. I named this argument loss, and adpated the tests. But maybe you have a convention how sth like this would be handled.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[fabian-sp]

Regarding the second bullet in the original post, I think I have now solved this by using jax.lax.cond.

[fabianp]

Thanks a lot @fabian-sp ! This looks great. I'll make a more through review once I get the vanilla Polyak SGD working (#718) :-)

[fabianp]

hey @fabian-sp , so sorry for the huge delay on this .... Merging the Polyak step-size highlighted some subtle issues that ended all the way in fixing stuff in pytype ... anyway, we got it merged finally, and now we can focus on MoMo!

A few of high level before we do a detailed review:

1. Since you first coded this, @vroulet added some common tests to the contrib directory (https://github.com/google-deepmind/optax/blob/main/optax/contrib/_common_test.py). Please add your solver to this file (under _OPTIMIZERS_UNDER_TEST)

2. Please merge the files momo.py and momo_adam.py into a single file, and the same for the *_test.py files.

3. Please take a look at the implementation of polyak_sgd in https://github.com/google-deepmind/optax/blob/main/optax/_src/alias.py . I suspect (but could be wrong) that a similar structure that splits the computation of the update from that of the step-size would make sense here too.

Thanks for all the work! 🙏🏼

[fabian-sp]

Thanks @fabianp, I merged the files. Somehow, after updating to the latest main, if I run locally the test.sh I get the following error

ERROR: Cannot install optax 0.1.9.dev0 (from /home/schaipp/uni/other/optax/dist/optax-0.1.9.dev0.tar.gz) and optax 0.2.2.dev0 (from /home/schaipp/uni/other/optax/dist/optax-0.2.2.dev0.tar.gz) because these package versions have conflicting dependencies.

Do you know how to fix this? Maybe some old installation of the optax package, but I thought the tester script would install in a new environment anyhow?

[fabianp]

you might need to uninstall optax pip uninstall optax before running the tests

[fabianp]

but yeah, its strange ...

[fabian-sp]

pip uninstall didnt solve it. Deleting all files in dist/ did the trick though.

The Github action tests are failing now, not sure why but probably because of changes made in the rest of the package since then?

[fabian-sp]

Okay, so the issue was that Momo needs the loss function value in the update (like Polyak SGD). This seems to be incompatible with the common tests. So I removed it for now from there, and tests are passing locally.

Depending on how you solved this for Polyak SGD, we can do it the same for Momo. @fabianp @vroulet

[fabianp]

so we basically wrote an if/else for polyak_sgd: https://github.com/google-deepmind/optax/blob/main/optax/_src/alias_test.py

[fabian-sp]

Okay, from my side the changes you requested should be implemented. A minor question: for Polyak-SGD you call the loss value argument value, while for MoMo I called it loss. This is mostly because I am used to Pytorch - feel free to change this if you prefer value.

[fabianp]

yes, please use value so it's consistent with the rest of optimizers in optax

[fabianp]

Thanks again for the contribution! Some minor comments here

[fabianp]

how about calling this f_min as with polyak_sgd ? https://optax.readthedocs.io/en/latest/api/optimizers.html#optax.polyak_sgd

[fabian-sp]

I would prefer the term lower bound, as this is closer to how we describe this quantity in the paper (strictly speaking, you don't need the optimal value for deriving MoMo, but only a lower bound). But I see your point of consistent naming...

[fabianp]

ok, but let's at least name it lower_bound ? lb is not very descriptive

[fabian-sp]

sure, can do that. Just fyi that I now added also an option where the lower bounded is estimated on the fly, which then will make some variable names a bit lengthy :) but at least the function argument should have a descriptive name

[fabianp]

yes exactly, having descriptive names in the function signature is the most important. Its fine if private variables have more cyptic names.

[fabianp]

have you considered writing this optimizer as a chain using the existing scale_by_adam ? If possible it would probably result in a much shorter (and reusable) code.

[fabian-sp]

For computing the adaptive learning rate, we need to compute the Adam EMAs, and some other quantities based on them. So I thought, it would be best to avoid double computations to have all in one function. But I might be wrong here...

[fabianp]

No, I think what you say makes sense

[fabian-sp]

Somehow after changing the formatting suggestions of @fabianp , now one test fails (the injection test for momo_adam). It's a bit mysterious for me, maybe you can have a look.

I will also now implement the adaptive lower bound estimation, that we proposed in the paper.

[fabianp]

Thanks! For the injection test, you might want to try converting the elements of the state into arrays. I vaguely recall having similar problems and that solving the issue

[fabian-sp]

Thanks, I did this. But the problem was solved by using a loss value other than zero here:

optax/optax/contrib/_common_test.py

Line 117 in 97128c1

update_kwargs = {'value': jnp.array(1.)}

With zero, it takes no step (as the Polyak step size is zero), and so it compares numerically zero values against each other, which seems to fail.

[fabianp]

FYI there are some test failures triggered by the last jax release that are unrelated to this PR (so don't worry about that for now, we're working on fixing them in parallel)

[fabian-sp]

Anything left for me to do, or are you gonna merge after fixing the upstream bug?

[fabianp]

Can you update your branch from master? hopefully the tests will then run again

[fabianp]

hey Fabian, sorry for the delay on this one. I was about to merge yesterday, but I realized what seemed to me like duplicated code between test_momo.py and _test_common.py . In particular, it seems that test_momo is testing on the same parabola and rosenbrock function that _test_commo.py . Am I missing something?

[fabianp]

It's now merged without test_momo.py (it was throwing some errors on our internal tests). Let me know how important it was and we can add it back in a follow up PR

[fabian-sp]

Hi Fabian, thanks for the final checks. The test_momo.py indeed had the initial tests, but later I used the common test structure of _test_common.py. So deleting it should be fine.

Thanks again for the interest in MoMo and the final push! 📦