[Feature] Implement more Shampoo features

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'])