search

google-deepmind / optax

Optax is a gradient processing and optimization library for JAX.
https://optax.readthedocs.io
Apache License 2.0
1.56k stars 165 forks source link

Implement Schedule-Free Learning #910

Closed ameya98 closed 3 months ago

ameya98 commented 3 months ago

From https://github.com/facebookresearch/schedule_free/tree/main by @adefazio and others.

I have a barebones implementation ready at https://gist.github.com/ameya98/7f103501714f4d2fdc0cb793579648d9, and would be happy to merge into Optax.

fabianp commented 3 months ago

Thanks @ameya98. There's now an implementation by @nullstring in https://github.com/google-deepmind/optax/pull/911 . Feel free to provide feedback for that pull request!