cleaned up utils.optim and wrote compliant NAG optim

Author: Alexander Ororbia

ngclearn/utils/optim/adam.py

@@ -5,7 +5,7 @@
 from functools import partial
 
 
-def step_update(param, update, g1, g2, lr, beta1, beta2, time_step, eps):
+def step_update(param, update, g1, g2, eta, beta1, beta2, time_step, eps):
     """
     Runs one step of Adam over a set of parameters given updates.
     The dynamics for any set of parameters is as follows:
@@ -28,7 +28,7 @@ def step_update(param, update, g1, g2, lr, beta1, beta2, time_step, eps):
         g2: second moment factor/correction factor to use in parameter update
             (must be same shape as "update")
 
-        lr: global step size value to be applied to updates to parameters
+        eta: global step size value to be applied to updates to parameters
 
         beta1: 1st moment control factor
 
@@ -45,7 +45,7 @@ def step_update(param, update, g1, g2, lr, beta1, beta2, time_step, eps):
     _g2 = beta2 * g2 + (1. - beta2) * jnp.square(update)
     g1_unb = _g1 / (1. - jnp.power(beta1, time_step))
     g2_unb = _g2 / (1. - jnp.power(beta2, time_step))
-    _param = param - lr * g1_unb/(jnp.sqrt(g2_unb) + eps)
+    _param = param - eta * g1_unb/(jnp.sqrt(g2_unb) + eps)
     return _param, _g1, _g2
 
 @jit

ngclearn/utils/optim/nag.py

@@ -6,7 +6,7 @@
 import time
 
 
-def step_update(param, update, phi_old, lr, mu, time_step):
+def step_update(param, update, phi_old, eta, mu, time_step):
     """
     Runs one step of Nesterov's accelerated gradient (NAG) over a set of parameters given updates.
     The dynamics for any set of parameters is as follows:
@@ -22,7 +22,7 @@ def step_update(param, update, phi_old, lr, mu, time_step):
 
         phi_old: previous friction/momentum parameter
 
-        lr: global step size value to be applied to updates to parameters
+        eta: global step size value to be applied to updates to parameters
 
         mu: friction/momentum control factor
 
@@ -31,7 +31,7 @@ def step_update(param, update, phi_old, lr, mu, time_step):
     Returns:
         adjusted parameter tensor (same shape as "param"), adjusted momentum/friction variable
     """
-    phi = param - update * lr ## do a phantom gradient adjustment step
+    phi = param - update * eta ## do a phantom gradient adjustment step
     _param = phi + (phi - phi_old) * (mu * (time_step > 1.)) ## NAG-step
     _phi_old = phi
     return _param, _phi_old

ngclearn/utils/optim/sgd.py

@@ -1,16 +1,16 @@
 from jax import jit, numpy as jnp
 
-def step_update(param, update, lr):
+def step_update(param, update, eta):
     """
     Runs one step of SGD over a set of parameters given updates.
 
     Args:
-        lr: global step size to apply when adjusting parameters
+        eta: global step size to apply when adjusting parameters
 
     Returns:
         adjusted parameter tensor (same shape as "param")
     """
-    _param = param - lr * update
+    _param = param - update * eta
     return _param
 
 @jit