Repository: PennyLaneAI/pennylane

Issue:
[BUG] Bugs in optimizer implementation and documentation
### Expected behavior

The optimizer tests refactor #1919 revealed issues in two of the optimizer classes.

**Adagrad optimizer**:
- The `apply_grad` method which applies a gradient step to a set of trainable parameters does not do so correctly, due to a missing index increment (`trained_index` in [adagrad.py#L76](https://github.com/PennyLaneAI/pennylane/blob/d705f0e3cc51bd876c8b910cd25b85bfc08694b0/pennylane/optimize/adagrad.py#L76)).

**Adam optimizer**: 
- The docstring for this class contains some error in the optimizer update rules, which are based on the paper [arXiv:1412.6980](https://arxiv.org/abs/1412.6980). The update rules are computed using biased first and second moment terms. A bias correction can then be applied to either the moments themselves or absorbed into a time-dependent step size. The current update rules do both of these, applying the correction twice.

- The optimizer implementation computes the time-dependent step size with `sqrt(1-b2^t)/(1-b1^t)` instead of `sqrt(1-b2)^t/(1-b1)^t` ([adam.py#L85](https://github.com/PennyLaneAI/pennylane/blob/8c6c4cc701f1eb00cc77aad1cbb6422c7dcf9f96/pennylane/optimize/adam.py#L85)). This seems to make no difference numerically in the tests but did stick out.

- The optimizer implementation omits factors of `(1-b1)` and `(1-b2)` in the first calculation of the first and second moment, which should be there since the bias correction (divisor) is absorbed into the step size ([adam.py#L126](https://github.com/PennyLaneAI/pennylane/blob/8c6c4cc701f1eb00cc77aad1cbb6422c7dcf9f96/pennylane/optimize/adam.py#L126) and [adam.py#L135](https://github.com/PennyLaneAI/pennylane/blob/8c6c4cc701f1eb00cc77aad1cbb6422c7dcf9f96/pennylane/optimize/adam.py#L135)). This seems erroneous although in the grand scheme of the optimization it probably doesn't matter.


### Actual behavior

see above

### Additional information

_No response_

### Source code

_No response_

### Tracebacks

_No response_

### System information

```shell
not needed
```


### 

- [X] I have searched exisisting GitHub issues to make sure the issue does not already exist.


Discussion:

Resolving patch:
diff --git a/pennylane/optimize/adagrad.py b/pennylane/optimize/adagrad.py
--- a/pennylane/optimize/adagrad.py
+++ b/pennylane/optimize/adagrad.py
@@ -74,6 +74,7 @@ def apply_grad(self, grad, args):
             if getattr(arg, "requires_grad", True):
                 x_flat = _flatten(arg)
                 grad_flat = list(_flatten(grad[trained_index]))
+                trained_index += 1
 
                 self._update_accumulation(index, grad_flat)
 
diff --git a/pennylane/optimize/adam.py b/pennylane/optimize/adam.py
--- a/pennylane/optimize/adam.py
+++ b/pennylane/optimize/adam.py
@@ -29,17 +29,17 @@ class AdamOptimizer(GradientDescentOptimizer):
     .. math::
         x^{(t+1)} = x^{(t)} - \eta^{(t+1)} \frac{a^{(t+1)}}{\sqrt{b^{(t+1)}} + \epsilon },
 
-    where the update rules for the three values are given by
+    where the update rules for the two moments are given by
 
     .. math::
-        a^{(t+1)} &= \frac{\beta_1 a^{(t)} + (1-\beta_1)\nabla f(x^{(t)})}{(1- \beta_1)},\\
-        b^{(t+1)} &= \frac{\beta_2 b^{(t)} + (1-\beta_2) ( \nabla f(x^{(t)}))^{\odot 2} }{(1- \beta_2)},\\
-        \eta^{(t+1)} &= \eta^{(t)} \frac{\sqrt{(1-\beta_2)}}{(1-\beta_1)}.
+        a^{(t+1)} &= \beta_1 a^{(t)} + (1-\beta_1) \nabla f(x^{(t)}),\\
+        b^{(t+1)} &= \beta_2 b^{(t)} + (1-\beta_2) (\nabla f(x^{(t)}))^{\odot 2},\\
+        \eta^{(t+1)} &= \eta \frac{\sqrt{(1-\beta_2^{t+1})}}{(1-\beta_1^{t+1})}.
 
     Above, :math:`( \nabla f(x^{(t-1)}))^{\odot 2}` denotes the element-wise square operation,
-    which means that each element in the gradient is multiplied by itself. The
-    hyperparameters :math:`\beta_1` and :math:`\beta_2` can also be step-dependent.
-    Initially, the first and second moment are zero.
+    which means that each element in the gradient is multiplied by itself. The hyperparameters
+    :math:`\beta_1` and :math:`\beta_2` can also be step-dependent. Initially, the first and
+    second moment are zero.
 
     The shift :math:`\epsilon` avoids division by zero.
 
@@ -123,7 +123,7 @@ def _update_moments(self, index, grad_flat):
         """
         # update first moment
         if self.accumulation["fm"][index] is None:
-            self.accumulation["fm"][index] = grad_flat
+            self.accumulation["fm"][index] = [(1 - self.beta1) * g for g in grad_flat]
         else:
             self.accumulation["fm"][index] = [
                 self.beta1 * f + (1 - self.beta1) * g
@@ -132,11 +132,11 @@ def _update_moments(self, index, grad_flat):
 
         # update second moment
         if self.accumulation["sm"][index] is None:
-            self.accumulation["sm"][index] = [g * g for g in grad_flat]
+            self.accumulation["sm"][index] = [(1 - self.beta2) * g * g for g in grad_flat]
         else:
             self.accumulation["sm"][index] = [
-                self.beta2 * f + (1 - self.beta2) * g * g
-                for f, g in zip(self.accumulation["sm"][index], grad_flat)
+                self.beta2 * s + (1 - self.beta2) * g * g
+                for s, g in zip(self.accumulation["sm"][index], grad_flat)
             ]
 
     def reset(self):


Test patch:
diff --git a/tests/optimize/test_optimize.py b/tests/optimize/test_optimize.py
--- a/tests/optimize/test_optimize.py
+++ b/tests/optimize/test_optimize.py
@@ -508,8 +508,8 @@ def test_adam_optimizer_univar(self, x_start, bunch, tol):
 
             x_onestep = bunch.adam_opt.step(f, x_start)
             adapted_stepsize = stepsize * np.sqrt(1 - delta) / (1 - gamma)
-            firstmoment = gradf(x_start)[0]
-            secondmoment = gradf(x_start)[0] * gradf(x_start)[0]
+            firstmoment = (1 - gamma) * gradf(x_start)[0]
+            secondmoment = (1 - delta) * gradf(x_start)[0] * gradf(x_start)[0]
             x_onestep_target = x_start - adapted_stepsize * firstmoment / (
                 np.sqrt(secondmoment) + 1e-8
             )
@@ -517,10 +517,9 @@ def test_adam_optimizer_univar(self, x_start, bunch, tol):
 
             x_twosteps = bunch.adam_opt.step(f, x_onestep)
             adapted_stepsize = stepsize * np.sqrt(1 - delta ** 2) / (1 - gamma ** 2)
-            firstmoment = gamma * gradf(x_start)[0] + (1 - gamma) * gradf(x_onestep)[0]
+            firstmoment = gamma * firstmoment + (1 - gamma) * gradf(x_onestep)[0]
             secondmoment = (
-                delta * gradf(x_start)[0] * gradf(x_start)[0]
-                + (1 - delta) * gradf(x_onestep)[0] * gradf(x_onestep)[0]
+                delta * secondmoment + (1 - delta) * gradf(x_onestep)[0] * gradf(x_onestep)[0]
             )
             x_twosteps_target = x_onestep - adapted_stepsize * firstmoment / (
                 np.sqrt(secondmoment) + 1e-8
@@ -538,8 +537,8 @@ def test_adam_optimizer_multivar(self, bunch, tol):
                 x_vec = x_vals[jdx : jdx + 2]
                 x_onestep = bunch.adam_opt.step(f, x_vec)
                 adapted_stepsize = stepsize * np.sqrt(1 - delta) / (1 - gamma)
-                firstmoment = gradf(x_vec)[0]
-                secondmoment = gradf(x_vec)[0] * gradf(x_vec)[0]
+                firstmoment = (1 - gamma) * gradf(x_vec)[0]
+                secondmoment = (1 - delta) * gradf(x_vec)[0] * gradf(x_vec)[0]
                 x_onestep_target = x_vec - adapted_stepsize * firstmoment / (
                     np.sqrt(secondmoment) + 1e-8
                 )
@@ -547,10 +546,9 @@ def test_adam_optimizer_multivar(self, bunch, tol):
 
                 x_twosteps = bunch.adam_opt.step(f, x_onestep)
                 adapted_stepsize = stepsize * np.sqrt(1 - delta ** 2) / (1 - gamma ** 2)
-                firstmoment = gamma * gradf(x_vec)[0] + (1 - gamma) * gradf(x_onestep)[0]
+                firstmoment = gamma * firstmoment + (1 - gamma) * gradf(x_onestep)[0]
                 secondmoment = (
-                    delta * gradf(x_vec)[0] * gradf(x_vec)[0]
-                    + (1 - delta) * gradf(x_onestep)[0] * gradf(x_onestep)[0]
+                    delta * secondmoment + (1 - delta) * gradf(x_onestep)[0] * gradf(x_onestep)[0]
                 )
                 x_twosteps_target = x_onestep - adapted_stepsize * firstmoment / (
                     np.sqrt(secondmoment) + 1e-8