Add inverse sqrt learning rate scheduler

[Sager611]

What does this PR do?

Adds the original invserse sqrt learning rate scheduler from Vaswani et al. (2017).

It is argued that this scheduler achieves the best performance when scaling ViTs on indefinite training times (Zhai et al. (2022)).

This PR ads a get_inverse_sqrt_schedule function and also updates the test in tests/optimization/test_optimization.py and the docs.

The implementation is adapted from:

• https://github.com/google-research/big_vision/blob/f071ce68852d56099437004fd70057597a95f6ef/big_vision/utils.py#L930

Timescale equals the no. of warmup steps by default as in:

• https://github.com/google-research/big_vision/blob/fd2d3bd2efc9d89ea959f16cd2f58ae8a495cd44/big_vision/configs/proj/clippo/train_clippo.py#L144
• https://github.com/google-research/big_vision/blob/6ff6d080d62c1f47e2e4eeb8b6474deb38dfe406/big_vision/configs/proj/scaling_laws/train_vit_g.py#L79

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[sgugger]

Thanks for adding this! I just have one comment.

[sgugger]

This should be a comment in the doc, the user reading the documentation won't really care about this.

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.