Repository: google-deepmind/optax

Issue:
`noisy_sgd` adds noise after scale by learning rate.
Thanks for the effort in this awesome library!

According to [these](https://github.com/deepmind/optax/blob/master/optax/_src/alias.py#L408#L435), it seems like the noise is added after scaling by the learning rate. That is, the noise is independent of the learning rate. I'm curious if this is an intentional design?


Discussion:
Good catch, this does not seem correct. My understanding is the noise is supposed to be added to the raw gradient, independent of the actual optimization method. That is, learning rate and any other scaling should take place after the noise is added to the gradient. Below is a snippet from the paper: https://arxiv.org/pdf/1511.06807.pdf

![image](https://user-images.githubusercontent.com/12474257/175794719-716714a2-1d7c-4b07-83bc-3d62defbc600.png)

The order has been this way since the library was originally published on GitHub (https://github.com/deepmind/optax/blame/d0ba208a300d68a3f4e380ed2a0b70848380988b/optax/_src/alias.py#L37-L39), so perhaps @mtthss has some insight? If there's no reason to keep this order, I think it makes sense to correct this.

I spoke with one of the authors and I can confirm that the noise should be added before multiplying by the learning rate.
@HanGuo97 do you want to put together a PR fixing this?
Unfortunately, I'm swamped with other responsibilities at the moment. I can come back to this later. In the meanwhile, contributions from others are more than welcome!
Shouldn't the fix just be this?

![Screen Shot 2022-06-29 at 20 42 52](https://user-images.githubusercontent.com/105969161/176512133-9f444cfe-0526-4950-8313-08bd90c04c23.png)

Not really sure how to validate this, so just ran the `test_optimization` test in `alias_test.py`. Here's the before & after:
<details>
<summary>
_setup_parabola params at every 1000th step
</summary>

#### Before

```
[-1. 10.  1.]
[0.7257005 0.4779135 1.02636  ]
[ 0.96638435 -0.8105099   0.98194563]
[ 1.0049064  -0.97684264  1.0158088 ]
[ 0.98075825 -0.99098873  1.0355433 ]
[ 0.99105674 -0.9999055   0.99992895]
[ 0.99034065 -1.0062125   0.99989396]
[ 1.0017453 -1.0287409  1.0216172]
[ 0.99269974 -0.9854698   0.98960614]
[ 0.97596675 -0.9877699   0.9836876 ]
```

#### After

```
[-1. 10.  1.]
[0.7298754 0.4857168 0.9999733]
[ 0.9635119 -0.7993221  1.0000178]
[ 0.99506235 -0.97289306  0.9999842 ]
[ 0.9993529 -0.9963453  0.9999641]
[ 0.99991876 -0.99950504  1.0000001 ]
[ 0.9999969  -0.99992734  0.99999946]
[ 0.9999966  -0.99996185  0.9999785 ]
[ 1.0000074 -1.0000128  1.0000106]
[ 1.0000243 -1.000012   1.000017 ]
```

</details>

<details>
<summary>
_setup_rosenbrock params at every 1000th step
</summary>

#### Before

```
[0. 0.]
[0.6903125 0.4818015]
[0.8406682  0.70478874]
[0.90429413 0.8187136 ]
[0.93457   0.8734052]
[0.9956608 0.9991537]
[0.9871882  0.97070897]
[0.9705165  0.94068545]
[0.9777622 0.9557112]
[0.97439975 0.9505268 ]
```

#### After

```
[0. 0.]
[0.6738693 0.4525301]
[0.8221863  0.67520136]
[0.8925657 0.7962139]
[0.9321801  0.86867684]
[0.9561389 0.9140138]
[0.9712825 0.9432765]
[0.98105425 0.96239185]
[0.9874232  0.97495407]
[0.99163187 0.9832988 ]
```

</details>
@atgctg that looks correct. With that change all the convergence tests still pass for me (rosenbrock and parabola). I don't think you need to verify this otherwise since the sub-components are well tested and the change is pretty simple.

Resolving patch:
diff --git a/optax/_src/alias.py b/optax/_src/alias.py
--- a/optax/_src/alias.py
+++ b/optax/_src/alias.py
@@ -430,8 +430,8 @@ def noisy_sgd(
     the corresponding `GradientTransformation`.
   """
   return combine.chain(
-      _scale_by_learning_rate(learning_rate),
       transform.add_noise(eta, gamma, seed),
+      _scale_by_learning_rate(learning_rate),
   )
 
 


Test patch: