Merge pull request #134 from kozistr/feature/roto-grad-optimizer

[Feature] Implement RotoGrad optimizer

Author: kozistr

README.rst

@@ -140,7 +140,9 @@ You can check the supported optimizers & lr schedulers.
 +--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 | SM3          | *Memory-Efficient Adaptive Optimization*                                                        | `github <https://github.com/google-research/google-research/tree/master/sm3>`__   | `https://arxiv.org/abs/1901.11150 <https://arxiv.org/abs/1901.11150>`__                       |
 +--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
-| AdaNorm`     | *Adaptive Gradient Norm Correction based Optimizer for CNNs*                                    | `github <https://github.com/shivram1987/AdaNorm>`__                               | `https://arxiv.org/abs/2210.06364 <https://arxiv.org/abs/2210.06364>`__                       |
+| AdaNorm      | *Adaptive Gradient Norm Correction based Optimizer for CNNs*                                    | `github <https://github.com/shivram1987/AdaNorm>`__                               | `https://arxiv.org/abs/2210.06364 <https://arxiv.org/abs/2210.06364>`__                       |
++--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
+| RotoGrad     | *Gradient Homogenization in Multitask Learning*                                                 | `github <https://github.com/adrianjav/rotograd>`__                                | `https://openreview.net/pdf?id=T8wHz4rnuGL <https://openreview.net/pdf?id=T8wHz4rnuGL>`__     |
 +--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 
 Useful Resources
@@ -351,6 +353,8 @@ Citations
 
 `AdaNorm <https://github.com/shivram1987/AdaNorm/tree/main#citation>`__
 
+`RotoGrad <https://github.com/adrianjav/rotograd#citing>`__
+
 Citation
 --------
 

docs/optimizer_api.rst

@@ -217,6 +217,14 @@ DAdaptSGD
 .. autoclass:: pytorch_optimizer.DAdaptSGD
     :members:
 
+.. _DAdaptAdan:
+
+DAdaptAdan
+----------
+
+.. autoclass:: pytorch_optimizer.DAdaptAdan
+    :members:
+
 .. _AdamS:
 
 AdamS
@@ -280,3 +288,11 @@ AdaNorm
 
 .. autoclass:: pytorch_optimizer.AdaNorm
     :members:
+
+.. _RotoGrad:
+
+RotoGrad
+--------
+
+.. autoclass:: pytorch_optimizer.RotoGrad
+    :members:

docs/util_api.rst

@@ -177,3 +177,11 @@ reduce_max_except_dim
 
 .. autoclass:: pytorch_optimizer.reduce_max_except_dim
     :members:
+
+.. _get_global_gradient_norm:
+
+get_global_gradient_norm
+------------------------
+
+.. autoclass:: pytorch_optimizer.get_global_gradient_norm
+    :members:

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytorch_optimizer"
-version = "2.6.1"
+version = "2.7.0"
 description = "optimizer & lr scheduler implementations in PyTorch with clean-code, strict types. Also, including useful optimization ideas."
 license = "Apache-2.0"
 authors = ["kozistr <[email protected]>"]
@@ -103,6 +103,7 @@ target-version = "py39"
 omit = [
     "./pytorch_optimizer/optimizer/gsam.py",
     "./pytorch_optimizer/optimizer/fp16.py",
+    "./pytorch_optimizer/optimizer/rotograd.py",
 ]
 
 [build-system]

pytorch_optimizer/__init__.py

@@ -28,7 +28,7 @@
 from pytorch_optimizer.optimizer.agc import agc
 from pytorch_optimizer.optimizer.alig import AliG
 from pytorch_optimizer.optimizer.apollo import Apollo
-from pytorch_optimizer.optimizer.dadapt import DAdaptAdaGrad, DAdaptAdam, DAdaptSGD
+from pytorch_optimizer.optimizer.dadapt import DAdaptAdaGrad, DAdaptAdam, DAdaptAdan, DAdaptSGD
 from pytorch_optimizer.optimizer.diffgrad import DiffGrad
 from pytorch_optimizer.optimizer.fp16 import DynamicLossScaler, SafeFP16Optimizer
 from pytorch_optimizer.optimizer.gc import centralize_gradient
@@ -45,6 +45,7 @@
 from pytorch_optimizer.optimizer.radam import RAdam
 from pytorch_optimizer.optimizer.ranger import Ranger
 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.sgdp import SGDP
 from pytorch_optimizer.optimizer.shampoo import ScalableShampoo, Shampoo
@@ -68,6 +69,7 @@
     clip_grad_norm,
     disable_running_stats,
     enable_running_stats,
+    get_global_gradient_norm,
     get_optimizer_parameters,
     normalize_gradient,
     reduce_max_except_dim,
@@ -96,6 +98,7 @@
     DAdaptAdaGrad,
     DAdaptAdam,
     DAdaptSGD,
+    DAdaptAdan,
     AdamS,
     AdaFactor,
     Apollo,

pytorch_optimizer/optimizer/adafactor.py

@@ -18,6 +18,7 @@ class AdaFactor(Optimizer, BaseOptimizer):
     :param decay_rate: float. coefficient used to compute running averages of square gradient.
     :param weight_decay: float. weight decay (L2 penalty).
     :param clip_threshold: float. threshold of root-mean-square of final gradient update.
+    :param amsgrad: bool. whether to use the AMSBound variant.
     :param scale_parameter: bool. if true, learning rate is scaled by root-mean-square of parameter.
     :param relative_step: bool. if true, time-dependent learning rate is computed instead of external learning rate.
     :param warmup_init: bool. time-dependent learning rate computation depends on whether warm-up initialization
@@ -34,6 +35,7 @@ def __init__(
         decay_rate: float = -0.8,
         weight_decay: float = 0.0,
         clip_threshold: float = 1.0,
+        amsgrad: bool = False,
         scale_parameter: bool = True,
         relative_step: bool = True,
         warmup_init: bool = False,
@@ -45,6 +47,7 @@ def __init__(
         self.decay_rate = decay_rate
         self.weight_decay = weight_decay
         self.clip_threshold = clip_threshold
+        self.amsgrad = amsgrad
         self.relative_step = relative_step
         self.eps1 = eps1
         self.eps2 = eps2
@@ -54,6 +57,7 @@ def __init__(
         defaults: DEFAULTS = {
             'lr': lr,
             'weight_decay': weight_decay,
+            'amsgrad': amsgrad,
             'scale_parameter': scale_parameter,
             'relative_step': relative_step,
             'warmup_init': warmup_init,
@@ -94,6 +98,9 @@ def reset(self):
                 else:
                     state['exp_avg_sq'] = torch.zeros_like(grad)
 
+                if group['amsgrad']:
+                    state['exp_avg_sq_hat'] = torch.zeros_like(grad)
+
                 state['RMS'] = 0.0
 
     def get_lr(
@@ -169,6 +176,9 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                     else:
                         state['exp_avg_sq'] = torch.zeros_like(grad)
 
+                    if group['amsgrad']:
+                        state['exp_avg_sq_hat'] = torch.zeros_like(grad)
+
                     state['RMS'] = 0.0
 
                 state['RMS'] = self.get_rms(p)
@@ -190,15 +200,18 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                     exp_avg_sq_row.mul_(beta2_t).add_(update.mean(dim=-1), alpha=1.0 - beta2_t)
                     exp_avg_sq_col.mul_(beta2_t).add_(update.mean(dim=-2), alpha=1.0 - beta2_t)
 
-                    self.approximate_sq_grad(exp_avg_sq_row, exp_avg_sq_col, output=update)
+                    self.approximate_sq_grad(exp_avg_sq_row, exp_avg_sq_col, update)
                 else:
                     exp_avg_sq = state['exp_avg_sq']
                     exp_avg_sq.mul_(beta2_t).add_(update, alpha=1.0 - beta2_t)
                     torch.rsqrt(exp_avg_sq, out=update)
 
-                update.mul_(grad)
+                if group['amsgrad']:
+                    exp_avg_sq_hat = state['exp_avg_sq_hat']
+                    torch.max(exp_avg_sq_hat, 1 / update, out=exp_avg_sq_hat)
+                    torch.rsqrt(exp_avg_sq_hat / beta2_t, out=update)
 
-                # TODO: implement AMSGrad
+                update.mul_(grad)
 
                 update.div_((self.get_rms(update) / self.clip_threshold).clamp_(min=1.0)).mul_(lr)
 

pytorch_optimizer/optimizer/adan.py

@@ -8,6 +8,7 @@
 from pytorch_optimizer.base.optimizer import BaseOptimizer
 from pytorch_optimizer.base.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS
 from pytorch_optimizer.optimizer.gc import centralize_gradient
+from pytorch_optimizer.optimizer.utils import get_global_gradient_norm
 
 
 class Adan(Optimizer, BaseOptimizer):
@@ -20,6 +21,8 @@ class Adan(Optimizer, BaseOptimizer):
     :param weight_decouple: bool. decoupled weight decay.
     :param max_grad_norm: float. max gradient norm to clip.
     :param use_gc: bool. use gradient centralization.
+    :param r: float. EMA factor. between 0.9 ~ 0.99 is preferred.
+    :param adanorm: bool. whether to use the AdaNorm variant.
     :param eps: float. term added to the denominator to improve numerical stability.
     """
 
@@ -32,6 +35,8 @@ def __init__(
         weight_decouple: bool = False,
         max_grad_norm: float = 0.0,
         use_gc: bool = False,
+        r: float = 0.95,
+        adanorm: bool = False,
         eps: float = 1e-8,
     ):
         self.lr = lr
@@ -49,8 +54,12 @@ def __init__(
             'weight_decay': weight_decay,
             'weight_decouple': weight_decouple,
             'max_grad_norm': max_grad_norm,
+            'adanorm': adanorm,
             'eps': eps,
         }
+        if adanorm:
+            defaults.update({'r': r})
+
         super().__init__(params, defaults)
 
     def validate_parameters(self):
@@ -71,25 +80,21 @@ def reset(self):
                 state = self.state[p]
 
                 state['exp_avg'] = torch.zeros_like(p)
+                state['exp_avg_sq'] = torch.zeros_like(p)
                 state['exp_avg_diff'] = torch.zeros_like(p)
-                state['exp_avg_nest'] = torch.zeros_like(p)
                 state['previous_grad'] = torch.zeros_like(p)
+                if group['adanorm']:
+                    state['exp_grad_norm'] = torch.zeros((1,), dtype=p.dtype, device=p.device)
 
     @torch.no_grad()
     def get_global_gradient_norm(self) -> Union[torch.Tensor, float]:
         if self.defaults['max_grad_norm'] == 0.0:
             return 1.0
 
-        global_grad_norm = torch.zeros(1, dtype=torch.float32, device=self.param_groups[0]['params'][0].device)
-
-        for group in self.param_groups:
-            for p in group['params']:
-                if p.grad is not None:
-                    global_grad_norm.add_(torch.linalg.norm(p.grad).pow(2))
-
-        global_grad_norm.sqrt_()
+        global_grad_norm = get_global_gradient_norm(self.param_groups, self.param_groups[0]['params'][0].device)
+        global_grad_norm.sqrt_().add_(self.eps)
 
-        return torch.clamp(self.defaults['max_grad_norm'] / (global_grad_norm + self.eps), max=1.0)
+        return torch.clamp(self.defaults['max_grad_norm'] / global_grad_norm, max=1.0)
 
     @torch.no_grad()
     def step(self, closure: CLOSURE = None) -> LOSS:
@@ -122,35 +127,50 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 state = self.state[p]
                 if len(state) == 0:
                     state['exp_avg'] = torch.zeros_like(p)
+                    state['exp_avg_sq'] = torch.zeros_like(p)
                     state['exp_avg_diff'] = torch.zeros_like(p)
-                    state['exp_avg_nest'] = torch.zeros_like(p)
-                    state['previous_grad'] = grad.clone()
+                    state['previous_grad'] = grad.clone().mul_(-clip_global_grad_norm)
+                    if group['adanorm']:
+                        state['exp_grad_norm'] = torch.zeros((1,), dtype=grad.dtype, device=grad.device)
 
                 grad.mul_(clip_global_grad_norm)
 
                 if self.use_gc:
                     grad = centralize_gradient(grad, gc_conv_only=False)
 
-                grad_diff = -state['previous_grad']
+                grad_diff = state['previous_grad']
                 grad_diff.add_(grad)
-                state['previous_grad'].copy_(grad)
 
-                update = grad + beta2 * grad_diff
+                s_grad = grad
+                if group['adanorm']:
+                    grad_norm = torch.linalg.norm(grad)
+
+                    exp_grad_norm = state['exp_grad_norm']
+                    exp_grad_norm.mul_(group['r']).add_(grad_norm, alpha=1.0 - group['r'])
 
-                exp_avg, exp_avg_diff, exp_avg_nest = state['exp_avg'], state['exp_avg_diff'], state['exp_avg_nest']
+                    if exp_grad_norm > grad_norm:
+                        s_grad *= exp_grad_norm / grad_norm
 
-                exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)
+                exp_avg, exp_avg_sq, exp_avg_diff = state['exp_avg'], state['exp_avg_sq'], state['exp_avg_diff']
+
+                exp_avg.mul_(beta1).add_(s_grad, alpha=1.0 - beta1)
                 exp_avg_diff.mul_(beta2).add_(grad_diff, alpha=1.0 - beta2)
-                exp_avg_nest.mul_(beta3).addcmul_(update, update, value=1.0 - beta3)
 
-                de_nom = (exp_avg_nest.sqrt_() / bias_correction3_sq).add_(self.eps)
-                perturb = (exp_avg / bias_correction1 + beta2 * exp_avg_diff / bias_correction2).div_(de_nom)
+                grad_diff.mul_(beta2).add_(grad)
+                exp_avg_sq.mul_(beta3).addcmul_(grad_diff, grad_diff, value=1.0 - beta3)
+
+                de_nom = exp_avg_sq.sqrt()
+                de_nom.div_(bias_correction3_sq).add_(group['eps'])
 
                 if group['weight_decouple']:
                     p.mul_(1.0 - group['lr'] * group['weight_decay'])
-                    p.add_(perturb, alpha=-group['lr'])
-                else:
-                    p.add_(perturb, alpha=-group['lr'])
+
+                p.addcdiv_(exp_avg, de_nom, value=-group['lr'] / bias_correction1)
+                p.addcdiv_(exp_avg_diff, de_nom, value=-group['lr'] * beta2 / bias_correction2)
+
+                if not group['weight_decouple']:
                     p.div_(1.0 + group['lr'] * group['weight_decay'])
 
+                state['previous_grad'].copy_(-grad)
+
         return loss

pytorch_optimizer/optimizer/alig.py

@@ -6,6 +6,7 @@
 from pytorch_optimizer.base.exception import NoClosureError, NoSparseGradientError
 from pytorch_optimizer.base.optimizer import BaseOptimizer
 from pytorch_optimizer.base.types import CLOSURE, DEFAULTS, LOSS, PARAMETERS
+from pytorch_optimizer.optimizer.utils import get_global_gradient_norm
 
 
 class AliG(Optimizer, BaseOptimizer):
@@ -60,14 +61,10 @@ def reset(self):
     @torch.no_grad()
     def compute_step_size(self, loss: float) -> float:
         r"""Compute step_size."""
-        global_grad_norm: float = 0
+        global_grad_norm = get_global_gradient_norm(self.param_groups, torch.device('cpu'))
+        global_grad_norm.add_(self.eps)
 
-        for group in self.param_groups:
-            for p in group['params']:
-                if p.grad is not None:
-                    global_grad_norm += p.grad.norm(2.0).pow(2).item()
-
-        return loss / (global_grad_norm + self.eps)
+        return loss / global_grad_norm.item()
 
     @torch.no_grad()
     def step(self, closure: CLOSURE = None) -> LOSS: