Added mathematical description to Noisy SGD

[hmludwig]

I added mathematical description to noisy SGD to improve clarity (see #757)

The results look like this:

[fabianp]

looks great to me, thanks!

[fabianp]

First line needs to end with ".", so I would leave the first line as it is, and add your description below

[fabianp]

and the line after should be blank, so something like:

r"""A variant of SGD with added noise.

Noisy SGD is ...
"""

[hmludwig]

Could you clarify if the requirement for the first line to end with "." is driven by stylistic guidelines or technical reasons? The list of optimizers at the beginning of the documention isn't effected and displays the first sentence of the docstring correctly.

[fabianp]

its a requirement of the python style: https://peps.python.org/pep-0257/#multi-line-docstrings

Internally we get an error when importing the code if the docstring doesn't follow this convention

[fabianp]

you're right though that this would be something worth adding to https://optax.readthedocs.io/en/latest/development.html

[vroulet]

Thanks for doing this!
As you are on it, could you change the description of gamma lines 1046-1047 of this file by

gamma: A parameter controlling the annealing of noise over time ``t``, the
      variance decays according to ``(1+t)**(-gamma)``.

Thank you again!

[vroulet]

Thanks for doing that!
Shouldn't it rather be the following?
u_t \leftarrow -\alpha_t (g_t + N(0, \sigma_t^2))

[hmludwig]

Yes, you're right. We first add the noise to the gradient and then scale by the learning rate. I will review my more carefully in the future.

[fabianp]

thanks @hmludwig ! Having accurate, reliable and helpful documentation is the most important part of a scientific project like this one. Your contribution makes a big difference 🙌🏼