Merge pull request #121 from kozistr/update/nero-eps

[Update] Add epsilon to stablize the optimizng

Author: kozistr

.github/workflows/ci.yml

@@ -35,7 +35,6 @@ jobs:
       run: make check
     - name: Check test
       env:
-        LD_PRELOAD: /usr/lib/x86_64-linux-gnu/libtcmalloc.so.4
         PYTHONDONTWRITEBYTECODE: 1
       run: make test
     - name: Check codecov

pytorch_optimizer/optimizer/lookahead.py

@@ -69,7 +69,8 @@ def update(self, group: Dict):
                     param_state['slow_mom'] = torch.zeros_like(fast)
 
             slow = param_state['slow_param']
-            slow += (fast - slow) * self.alpha
+            slow.add_(fast - slow, alpha=self.alpha)
+
             fast.copy_(slow)
 
             if 'momentum_buffer' not in self.optimizer.state[fast]:
@@ -98,30 +99,21 @@ def step(self, closure: CLOSURE = None) -> LOSS:
         return loss
 
     def state_dict(self) -> STATE:
-        fast_state_dict: STATE = self.optimizer.state_dict()
-        fast_state = fast_state_dict['state']
-        param_groups = fast_state_dict['param_groups']
-
+        fast_state: STATE = self.optimizer.state_dict()
         slow_state: STATE = {(id(k) if isinstance(k, torch.Tensor) else k): v for k, v in self.state.items()}
 
         return {
-            'fast_state': fast_state,
+            'fast_state': fast_state['state'],
             'slow_state': slow_state,
-            'param_groups': param_groups,
+            'param_groups': fast_state['param_groups'],
         }
 
-    def load_state_dict(self, state_dict: STATE):
-        slow_state_dict: STATE = {
-            'state': state_dict['slow_state'],
-            'param_groups': state_dict['param_groups'],
-        }
-        fast_state_dict: STATE = {
-            'state': state_dict['fast_state'],
-            'param_groups': state_dict['param_groups'],
-        }
-        super().load_state_dict(slow_state_dict)
+    def load_state_dict(self, state: STATE):
+        slow_state: STATE = {'state': state['slow_state'], 'param_groups': state['param_groups']}
+        fast_state: STATE = {'state': state['fast_state'], 'param_groups': state['param_groups']}
+        super().load_state_dict(slow_state)
 
-        self.optimizer.load_state_dict(fast_state_dict)
+        self.optimizer.load_state_dict(fast_state)
         self.fast_state = self.optimizer.state
 
     def add_param_group(self, param_group):

pytorch_optimizer/optimizer/nero.py

@@ -14,11 +14,15 @@ class Nero(Optimizer, BaseOptimizer):
     :param lr: float. learning rate.
     :param beta: float. coefficients used for computing running averages of gradient and the squared hessian trace.
     :param constraints: bool.
+    :param eps: float. term added to the denominator to improve numerical stability.
     """
 
-    def __init__(self, params: PARAMETERS, lr: float = 0.01, beta: float = 0.999, constraints: bool = True):
+    def __init__(
+        self, params: PARAMETERS, lr: float = 0.01, beta: float = 0.999, constraints: bool = True, eps: float = 1e-8
+    ):
         self.lr = lr
         self.beta = beta
+        self.eps = eps
 
         self.validate_parameters()
 
@@ -28,6 +32,7 @@ def __init__(self, params: PARAMETERS, lr: float = 0.01, beta: float = 0.999, co
     def validate_parameters(self):
         self.validate_learning_rate(self.lr)
         self.validate_beta(self.beta)
+        self.validate_epsilon(self.eps)
 
     def __str__(self) -> str:
         return 'Nero'
@@ -38,7 +43,7 @@ def reset(self):
             for p in group['params']:
                 if group['constraints'] and p.dim() > 1:
                     p.sub_(neuron_mean(p))
-                    p.div_(neuron_norm(p))
+                    p.div_(neuron_norm(p) + self.eps)
 
                 state = self.state[p]
 
@@ -69,7 +74,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 if len(state) == 0:
                     if group['constraints'] and p.dim() > 1:
                         p.sub_(neuron_mean(p))
-                        p.div_(neuron_norm(p))
+                        p.div_(neuron_norm(p) + self.eps)
 
                     state['step'] = 0
                     state['exp_avg_sq'] = torch.zeros_like(neuron_norm(p))
@@ -79,16 +84,20 @@ def step(self, closure: CLOSURE = None) -> LOSS:
 
                 state['step'] += 1
 
+                grad_norm = neuron_norm(grad)
+
+                exp_avg_sq = state['exp_avg_sq']
+                exp_avg_sq.mul_(self.beta).addcmul_(grad_norm, grad_norm, value=1.0 - self.beta)
+
                 bias_correction: float = 1.0 - self.beta ** state['step']
-                state['exp_avg_sq'] = self.beta * state['exp_avg_sq'] + (1.0 - self.beta) * neuron_norm(grad) ** 2
 
-                grad_normed = grad / (state['exp_avg_sq'] / bias_correction).sqrt()
-                grad_normed[torch.isnan(grad_normed)] = 0.0
+                grad_normed = grad / ((exp_avg_sq / bias_correction).sqrt() + self.eps)
+                torch.nan_to_num(grad_normed, nan=0.0, out=grad_normed)
 
                 p.sub_(group['lr'] * state['scale'] * grad_normed)
 
                 if group['constraints'] and p.dim() > 1:
                     p.sub_(neuron_mean(p))
-                    p.div_(neuron_norm(p))
+                    p.div_(neuron_norm(p) + self.eps)
 
         return loss

pytorch_optimizer/optimizer/utils.py

@@ -199,21 +199,19 @@ def neuron_norm(x: torch.Tensor) -> torch.Tensor:
     if x.dim() <= 1:
         return x.abs()
 
-    view_shape = [x.shape[0]] + [1] * (x.dim() - 1)
-    x = x.view(x.shape[0], -1)
+    view_shape: List[int] = [x.shape[0]] + [1] * (x.dim() - 1)
 
-    return x.norm(dim=1).view(*view_shape)
+    return channel_view(x).norm(dim=1).view(*view_shape)
 
 
 def neuron_mean(x: torch.Tensor) -> torch.Tensor:
     r"""Get mean of the tensor."""
     if x.dim() <= 1:
         raise ValueError('[-] neuron_mean not defined on 1D tensors.')
 
-    view_shape = [x.shape[0]] + [1] * (x.dim() - 1)
-    x = x.view(x.shape[0], -1)
+    view_shape: List[int] = [x.shape[0]] + [1] * (x.dim() - 1)
 
-    return x.mean(dim=1).view(*view_shape)
+    return channel_view(x).mean(dim=1).view(*view_shape)
 
 
 def disable_running_stats(model):

tests/test_gradients.py

@@ -39,8 +39,6 @@ def test_sparse_not_supported(no_sparse_optimizer):
     opt = load_optimizer(optimizer=no_sparse_optimizer)
     optimizer = opt([param], num_iterations=1) if no_sparse_optimizer == 'ranger21' else opt([param])
 
-    optimizer.zero_grad()
-
     with pytest.raises(NoSparseGradientError):
         optimizer.step(lambda: 0.1)
 

tests/test_optimizer_parameters.py

@@ -25,7 +25,7 @@ def test_learning_rate(optimizer_name):
 
 @pytest.mark.parametrize('optimizer_name', VALID_OPTIMIZER_NAMES)
 def test_epsilon(optimizer_name):
-    if optimizer_name in ('nero', 'shampoo', 'scalableshampoo', 'dadaptsgd', 'adafactor', 'lion'):
+    if optimizer_name in ('shampoo', 'scalableshampoo', 'dadaptsgd', 'adafactor', 'lion'):
         pytest.skip(f'skip {optimizer_name} optimizer')
 
     optimizer = load_optimizer(optimizer_name)