Merge pull request #99 from kozistr/update/shampoo-optimizer

[Feature] Implement more Shampoo features

Author: kozistr

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytorch_optimizer"
-version = "2.3.0"
+version = "2.3.1"
 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]>"]

pytorch_optimizer/optimizer/shampoo.py

@@ -3,8 +3,17 @@
 
 from pytorch_optimizer.base.exception import NoSparseGradientError
 from pytorch_optimizer.base.optimizer import BaseOptimizer
-from pytorch_optimizer.base.types import CLOSURE, DEFAULTS, LOSS, PARAMETERS
-from pytorch_optimizer.optimizer.shampoo_utils import AdagradGraft, Graft, LayerWiseGrafting, PreConditioner, SGDGraft
+from pytorch_optimizer.base.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS
+from pytorch_optimizer.optimizer.shampoo_utils import (
+    AdagradGraft,
+    Graft,
+    LayerWiseGrafting,
+    PreConditioner,
+    PreConditionerType,
+    RMSPropGraft,
+    SGDGraft,
+    SQRTNGraft,
+)
 
 
 class Shampoo(Optimizer, BaseOptimizer):
@@ -14,9 +23,11 @@ class Shampoo(Optimizer, BaseOptimizer):
 
     :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
     :param lr: float. learning rate.
-    :param momentum: float. momentum.
-    :param beta2: float. beta2.
+    :param betas: BETAS. beta1, beta2.
+    :param moving_average_for_momentum: bool. perform moving_average for momentum (beta1).
     :param weight_decay: float. weight decay (L2 penalty).
+    :param decoupled_weight_decay: bool. use decoupled weight_decay.
+    :param decoupled_learning_rate: bool. use decoupled lr, otherwise couple it w/ preconditioned gradient.
     :param inverse_exponent_override: int. fixed exponent for pre-conditioner, if > 0.
     :param start_preconditioning_step: int.
     :param preconditioning_compute_steps: int. performance tuning params for controlling memory and compute
@@ -28,7 +39,8 @@ class Shampoo(Optimizer, BaseOptimizer):
     :param shape_interpretation: bool. Automatic shape interpretation (for eg: [4, 3, 1024, 512] would
         result in 12 x [1024, 512] L and R statistics. Disabled by default which results in Shampoo constructing
         statistics [4, 4], [3, 3], [1024, 1024], [512, 512].
-    :param graft_type: bool. Type of grafting (SGD or AdaGrad).
+    :param graft_type: int. type of grafting (SGD or AdaGrad or RMSProp or SQRT_N or None).
+    :param pre_conditioner_type: int. type of pre-conditioner.
     :param nesterov: bool. Nesterov momentum.
     :param diagonal_eps: float. term added to the denominator to improve numerical stability.
     :param matrix_eps: float. term added to the denominator to improve numerical stability.
@@ -38,31 +50,37 @@ def __init__(
         self,
         params: PARAMETERS,
         lr: float = 1e-3,
-        momentum: float = 0.0,
-        beta2: float = 1.0,
+        betas: BETAS = (0.9, 0.999),
+        moving_average_for_momentum: bool = False,
         weight_decay: float = 0.0,
+        decoupled_weight_decay: bool = False,
+        decoupled_learning_rate: bool = True,
         inverse_exponent_override: int = 0,
-        start_preconditioning_step: int = 1,
+        start_preconditioning_step: int = 5,
         preconditioning_compute_steps: int = 1,
         statistics_compute_steps: int = 1,
         block_size: int = 128,
         shape_interpretation: bool = True,
         graft_type: int = LayerWiseGrafting.SGD,
+        pre_conditioner_type: int = PreConditionerType.ALL,
         nesterov: bool = True,
-        diagonal_eps: float = 1e-6,
-        matrix_eps: float = 1e-12,
+        diagonal_eps: float = 1e-10,
+        matrix_eps: float = 1e-6,
     ):
         self.lr = lr
-        self.momentum = momentum
-        self.beta2 = beta2
+        self.betas = betas
+        self.moving_average_for_momentum = moving_average_for_momentum
         self.weight_decay = weight_decay
+        self.decoupled_weight_decay = decoupled_weight_decay
+        self.decoupled_learning_rate = decoupled_learning_rate
         self.inverse_exponent_override = inverse_exponent_override
         self.start_preconditioning_step = start_preconditioning_step
         self.preconditioning_compute_steps = preconditioning_compute_steps
         self.statistics_compute_steps = statistics_compute_steps
         self.block_size = block_size
         self.shape_interpretation = shape_interpretation
         self.graft_type = graft_type
+        self.pre_conditioner_type = pre_conditioner_type
         self.nesterov = nesterov
         self.diagonal_eps = diagonal_eps
         self.matrix_eps = matrix_eps
@@ -71,14 +89,14 @@ def __init__(
 
         defaults: DEFAULTS = {
             'lr': lr,
-            'momentum': momentum,
+            'betas': betas,
             'weight_decay': weight_decay,
         }
         super().__init__(params, defaults)
 
     def validate_parameters(self):
         self.validate_learning_rate(self.lr)
-        self.validate_momentum(self.momentum)
+        self.validate_betas(self.betas)
         self.validate_weight_decay(self.weight_decay)
         self.validate_update_frequency(self.start_preconditioning_step)
         self.validate_update_frequency(self.statistics_compute_steps)
@@ -100,16 +118,21 @@ def reset(self):
                 state['momentum'] = torch.zeros_like(p)
                 state['pre_conditioner'] = PreConditioner(
                     p,
-                    self.beta2,
+                    group['betas'][1],  # beta2
                     self.inverse_exponent_override,
                     self.block_size,
                     self.shape_interpretation,
                     self.matrix_eps,
+                    self.pre_conditioner_type,
                 )
                 if self.graft_type == LayerWiseGrafting.ADAGRAD:
                     state['graft'] = AdagradGraft(p, self.diagonal_eps)
+                elif self.graft_type == LayerWiseGrafting.RMSPROP:
+                    state['graft'] = RMSPropGraft(p, self.diagonal_eps)
                 elif self.graft_type == LayerWiseGrafting.SGD:
                     state['graft'] = SGDGraft(p)
+                elif self.graft_type == LayerWiseGrafting.SQRTN:
+                    state['graft'] = SQRTNGraft(p)
                 else:
                     state['graft'] = Graft(p)
 
@@ -121,6 +144,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 loss = closure()
 
         for group in self.param_groups:
+            beta1, beta2 = group['betas']
             for p in group['params']:
                 if p.grad is None:
                     continue
@@ -135,48 +159,59 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                     state['momentum'] = torch.zeros_like(p)
                     state['pre_conditioner'] = PreConditioner(
                         p,
-                        self.beta2,
+                        beta2,
                         self.inverse_exponent_override,
                         self.block_size,
                         self.shape_interpretation,
                         self.matrix_eps,
+                        self.pre_conditioner_type,
                     )
                     if self.graft_type == LayerWiseGrafting.ADAGRAD:
                         state['graft'] = AdagradGraft(p, self.diagonal_eps)
+                    elif self.graft_type == LayerWiseGrafting.RMSPROP:
+                        state['graft'] = RMSPropGraft(p, self.diagonal_eps)
                     elif self.graft_type == LayerWiseGrafting.SGD:
                         state['graft'] = SGDGraft(p)
+                    elif self.graft_type == LayerWiseGrafting.SQRTN:
+                        state['graft'] = SQRTNGraft(p)
                     else:
                         state['graft'] = Graft(p)
 
                 state['step'] += 1
                 pre_conditioner, graft = state['pre_conditioner'], state['graft']
 
                 # gather statistics, compute pre-conditioners
-                graft.add_statistics(grad)
+                graft.add_statistics(grad, beta2)
                 if state['step'] % self.statistics_compute_steps == 0:
                     pre_conditioner.add_statistics(grad)
                 if state['step'] % self.preconditioning_compute_steps == 0:
                     pre_conditioner.compute_pre_conditioners()
 
                 # pre-condition gradients
-                graft_grad: torch.Tensor = graft.precondition_gradient(grad)
+                pre_conditioner_multiplier: float = group['lr'] if not self.decoupled_learning_rate else 1.0
+                graft_grad: torch.Tensor = graft.precondition_gradient(grad * pre_conditioner_multiplier)
                 shampoo_grad: torch.Tensor = grad
                 if state['step'] >= self.start_preconditioning_step:
                     shampoo_grad = pre_conditioner.preconditioned_grad(grad)
 
                 # grafting
                 graft_norm = torch.norm(graft_grad)
                 shampoo_norm = torch.norm(shampoo_grad)
-                shampoo_grad.mul_(graft_norm / (shampoo_norm + 1e-16))
+                if self.graft_type != LayerWiseGrafting.NONE:
+                    shampoo_grad.mul_(graft_norm / (shampoo_norm + 1e-16))
 
                 # apply weight decay (adam style)
                 if group['weight_decay'] > 0.0:
-                    shampoo_grad.add_(p, alpha=group['weight_decay'])
-                    graft_grad.add_(p, alpha=group['weight_decay'])
+                    if not self.decoupled_weight_decay:
+                        shampoo_grad.add_(p, alpha=group['weight_decay'])
+                        graft_grad.add_(p, alpha=group['weight_decay'])
+                    else:
+                        shampoo_grad.mul_(1.0 - group['lr'] * group['weight_decay'])
+                        graft_grad.mul_(1.0 - group['lr'] * group['weight_decay'])
 
                 # Momentum and Nesterov momentum, if needed
-                state['momentum'].mul_(group['momentum']).add_(shampoo_grad)
-                graft_momentum = graft.update_momentum(grad, group['momentum'])
+                state['momentum'].mul_(beta1).add_(shampoo_grad)
+                graft_momentum = graft.update_momentum(grad, beta1)
 
                 if state['step'] >= self.start_preconditioning_step:
                     momentum_update = state['momentum']
@@ -186,7 +221,10 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                     wd_update = graft_grad
 
                 if self.nesterov:
-                    momentum_update.mul_(group['momentum']).add_(wd_update)
+                    w: float = (1.0 - beta1) if self.moving_average_for_momentum else 1.0
+                    wd_update.mul_(w)
+
+                    momentum_update.mul_(beta1).add_(wd_update)
 
                 p.add_(momentum_update, alpha=-group['lr'])