Add SignSGD, allow heesian approximation distribution selection, fix adahessian instability

Author: ferris

pytorch_optimizer/__init__.py

@@ -62,7 +62,7 @@
 from pytorch_optimizer.optimizer.ranger21 import Ranger21
 from pytorch_optimizer.optimizer.rotograd import RotoGrad
 from pytorch_optimizer.optimizer.sam import SAM
-from pytorch_optimizer.optimizer.sgd import ASGD, SGDW, AccSGD
+from pytorch_optimizer.optimizer.sgd import ASGD, SGDW, AccSGD, SignSGD
 from pytorch_optimizer.optimizer.sgdp import SGDP
 from pytorch_optimizer.optimizer.shampoo import ScalableShampoo, Shampoo
 from pytorch_optimizer.optimizer.shampoo_utils import (
@@ -150,7 +150,8 @@
     AdaDelta,
     Amos,
     AdaHessian,
-    SophiaH
+    SophiaH,
+    SignSGD
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
 

pytorch_optimizer/base/optimizer.py

@@ -5,17 +5,20 @@
 import torch
 
 from pytorch_optimizer.base.exception import NegativeLRError, NegativeStepError
-from pytorch_optimizer.base.types import BETAS
+from pytorch_optimizer.base.types import BETAS, HUTCHINSON_G
 
 
 class BaseOptimizer(ABC):
     r"""Base optimizer class."""
 
     @torch.no_grad()
-    def compute_hutchinson_hessian(self, nsamples: int = 1, pre_zero=True, alpha=1.0):
+    def compute_hutchinson_hessian(self, nsamples: int = 1, pre_zero=True, alpha=1.0, distribution: HUTCHINSON_G = 'gaussian'):
         """
         Hutchinsons approximate hessian, added to the state under key 'hessian'
         """
+        if distribution not in ['gaussian', 'rademacher']:
+            raise NotImplementedError(f"Hessian with distribution {distribution} is not implemented")
+
         params = []
         for group in self.param_groups:
             for p in group['params']:
@@ -34,14 +37,17 @@ def compute_hutchinson_hessian(self, nsamples: int = 1, pre_zero=True, alpha=1.0
         grads = [p.grad for p in params]
 
         for i in range(nsamples):
-            # Gaussian N(0,Id)
-            zs = [torch.randn(p.size(), device=p.device) for p in params]
-            # Rademacher distribution {-1.0, 1.0}
-            # zs = [torch.randint(0, 2, p.size(), device=p.device) * 2.0 - 1.0 for p in params]
+            if distribution == 'gaussian':
+                # Gaussian N(0,Id)
+                zs = [torch.randn(p.size(), device=p.device) for p in params]
+            elif distribution == 'rademacher':
+                # Rademacher distribution {-1.0, 1.0}
+                zs = [torch.randint(0, 2, p.size(), dtype=p.dtype, device=p.device) * 2.0 - 1.0 for p in params]
+
             h_zs = torch.autograd.grad(grads, params, grad_outputs=zs, retain_graph=i < nsamples - 1)
             for h_z, z, p in zip(h_zs, zs, params):
                 # approximate the expected values of z*(H@z)
-                self.state[p]['hessian'].add_(h_z * z, alpha=1/nsamples * alpha)
+                self.state[p]['hessian'].add_(h_z * z, alpha=(1/nsamples) * alpha)
 
     @staticmethod
     def apply_weight_decay(

pytorch_optimizer/base/types.py

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Dict, Iterable, Optional, Tuple, Type, Union
+from typing import Any, Callable, Dict, Iterable, Optional, Tuple, Type, Union, Literal
 
 import torch
 from torch.optim import Optimizer
@@ -12,3 +12,5 @@
 STATE = Dict[str, Any]
 OPTIMIZER = Type[Optimizer]
 SCHEDULER = Type[_LRScheduler]
+
+HUTCHINSON_G = Literal['gaussian', 'rademacher']

pytorch_optimizer/optimizer/adahessian.py

@@ -3,7 +3,7 @@
 
 from pytorch_optimizer.base.exception import NoSparseGradientError
 from pytorch_optimizer.base.optimizer import BaseOptimizer
-from pytorch_optimizer.base.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS
+from pytorch_optimizer.base.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS, HUTCHINSON_G
 
 # Modified from https://github.com/davda54/ada-hessian/blob/master/ada_hessian.py (MIT David Samuel)
 
@@ -35,6 +35,7 @@ def __init__(self,
                  hessian_power: float = 1.0,
                  update_period: int = 1,
                  n_samples: int = 1,
+                 hessian_distribution: HUTCHINSON_G = 'rademacher',
                  eps: float = 1e-16):
 
         self.validate_learning_rate(lr)
@@ -64,7 +65,7 @@ def reset(self):
             for p in group['params']:
                 state = self.state[p]
                 state['exp_avg'] = torch.zeros_like(p)
-                state['exp_hessian_diag_sq'] = torch.zeros_like(p)
+                state['exp_hessian_diag_sq'] = state['hessian'].clone()
 
     @torch.no_grad()
     def step(self, closure: CLOSURE = None) -> LOSS:
@@ -88,8 +89,9 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 # State initialization
                 state = self.state[p]
                 if 'exp_avg' not in state:
-                    state['exp_avg'] = torch.zeros_like(p.data)  # Exponential moving average of gradient values
-                    state['exp_hessian_diag_sq'] = torch.zeros_like(p.data)  # Exponential moving average of Hessian diagonal square values
+                    state['exp_avg'] = torch.zeros_like(p.data)
+                    # NOTE: zeroing-out the hessian causes instability
+                    state['exp_hessian_diag_sq'] = state['hessian'].clone()
 
                 self.apply_weight_decay(
                     p=p,

pytorch_optimizer/optimizer/sgd.py

@@ -311,3 +311,75 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 p.add_(grad, alpha=-new_lr)
 
         return loss
+
+
+class SignSGD(Optimizer, BaseOptimizer):
+    r"""SignSGD: Compressed Optimisation for Non-Convex Problems
+
+    :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
+    :param lr: float. learning rate.
+    :param momentum: float. momentum factor (0.0=SignSGD, >0=Signum).
+    :param weight_decay: float. weight decay (L2 penalty).
+    :param weight_decouple: bool. the optimizer uses decoupled weight decay as in AdamW.
+    """
+
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 1e-3,
+        beta: float = 0.9,
+        weight_decay: float = 0.0,
+        weight_decouple: bool = True,
+    ):
+        self.validate_learning_rate(lr)
+        self.validate_range(beta, 'beta', 0.0, 1.0)
+        self.validate_non_negative(weight_decay, 'weight_decay')
+
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'beta': beta,
+            'weight_decay': weight_decay,
+            'weight_decouple': weight_decouple,
+        }
+
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def reset(self):
+        for group in self.param_groups:
+            for p in group['params']:
+                state = self.state[p]
+
+                if group['beta'] > 0.0:
+                    state['momentum_buffer'] = p.grad.clone()
+
+    @torch.no_grad()
+    def step(self, closure: CLOSURE = None) -> LOSS:
+        loss: LOSS = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            beta = group['beta']
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                if p.grad.is_sparse:
+                    raise NoSparseGradientError(str(self))
+
+                state = self.state[p]
+
+                if beta > 0.0:
+                    if len(state) == 0:
+                        state['momentum_buffer'] = p.grad.clone()
+
+                    buf = state['momentum_buffer']
+                    buf.mul_(beta).add_(p.grad, alpha=1.0 - beta)
+                else:
+                    buf = p.grad
+
+                p.add_(torch.sign(buf), alpha=-group['lr'])
+
+        return loss

pytorch_optimizer/optimizer/sophiah.py

@@ -3,7 +3,7 @@
 
 from pytorch_optimizer.base.exception import NoSparseGradientError
 from pytorch_optimizer.base.optimizer import BaseOptimizer
-from pytorch_optimizer.base.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS
+from pytorch_optimizer.base.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS, HUTCHINSON_G
 
 
 class SophiaH(Optimizer, BaseOptimizer):
@@ -33,6 +33,7 @@ def __init__(self,
                  p: float = 25.,
                  update_period: int = 10,
                  n_samples: int = 1,
+                 hessian_distribution: HUTCHINSON_G = 'gaussian',
                  eps: float = 1e-12):
 
         self.validate_learning_rate(lr)