feature: implement SimplifiedAdEMAMix optimizer

Author: kozistr

pytorch_optimizer/optimizer/ademamix.py

@@ -161,3 +161,137 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 p.add_(update, alpha=-step_size)
 
         return loss
+
+
+class SimplifiedAdEMAMix(BaseOptimizer):
+    r"""Connections between Schedule-Free Optimizers, AdEMAMix, and Accelerated SGD Variants.
+
+    :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
+    :param lr: float. learning rate.
+    :param betas: BETAS. coefficients used for computing running averages of gradient and the squared hessian trace.
+    :param alpha: float. coefficient for mixing the current gradient and EMA.
+    :param beta1_warmup: Optional[int]. number of warmup steps used to increase beta1.
+    :param min_beta1: float. minimum value of beta1 to start from.
+    :param weight_decay: float. weight decay (L2 penalty).
+    :param weight_decouple: bool. the optimizer uses decoupled weight decay as in AdamW.
+    :param fixed_decay: bool. fix weight decay.
+    :param eps: float. term added to the denominator to improve numerical stability.
+    """
+
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 1e-4,
+        betas: BETAS = (0.99, 0.95),
+        weight_decay: float = 0.0,
+        weight_decouple: bool = True,
+        fixed_decay: bool = False,
+        alpha: float = 0.0,
+        beta1_warmup: Optional[int] = None,
+        min_beta1: float = 0.9,
+        eps: float = 1e-8,
+        **kwargs,
+    ):
+        self.validate_learning_rate(lr)
+        self.validate_betas(betas)
+        self.validate_non_negative(alpha, 'alpha')
+        self.validate_non_negative(min_beta1, 'min_beta1')
+        self.validate_non_negative(weight_decay, 'weight_decay')
+        self.validate_non_negative(eps, 'eps')
+
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'betas': betas,
+            'alpha': alpha,
+            'beta1_warmup': beta1_warmup,
+            'min_beta1': min_beta1,
+            'weight_decay': weight_decay,
+            'weight_decouple': weight_decouple,
+            'fixed_decay': fixed_decay,
+            'eps': eps,
+        }
+
+        super().__init__(params, defaults)
+
+    def __str__(self) -> str:
+        return 'SimAdEMAMix'
+
+    @torch.no_grad()
+    def reset(self):
+        pass
+
+    @staticmethod
+    def linear_hl_warmup_scheduler(step: int, beta_end: float, beta_start: float = 0.0, warmup: int = 1) -> float:
+
+        def f(beta: float, eps: float = 1e-8) -> float:
+            return math.log(0.5) / math.log(beta + eps) - 1.0
+
+        def f_inv(t: float) -> float:
+            return math.pow(0.5, 1.0 / (t + 1))
+
+        if step < warmup:
+            a: float = step / float(warmup)
+            return f_inv((1.0 - a) * f(beta_start) + a * f(beta_end))
+
+        return beta_end
+
+    @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:
+            if 'step' in group:
+                group['step'] += 1
+            else:
+                group['step'] = 1
+
+            beta1, beta2 = group['betas']
+
+            if group['beta1_warmup']:
+                beta1 = self.linear_hl_warmup_scheduler(
+                    group['step'], beta_end=beta1, beta_start=group['min_beta1'], warmup=group['beta1_warmup']
+                )
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad
+                if grad.is_sparse:
+                    raise NoSparseGradientError(str(self))
+
+                state = self.state[p]
+
+                if len(state) == 0:
+                    state['exp_avg'] = torch.zeros_like(p)
+                    state['exp_avg_sq'] = torch.zeros_like(p)
+                    state['num_sum'] = 0.0
+                    state['den_sum'] = 0.0
+
+                self.apply_weight_decay(
+                    p=p,
+                    grad=grad,
+                    lr=group['lr'],
+                    weight_decay=group['weight_decay'],
+                    weight_decouple=group['weight_decouple'],
+                    fixed_decay=group['fixed_decay'],
+                )
+
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+
+                exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
+
+                state['num_sum'] = beta1 * state['num_sum'] + 1.0
+                state['den_sum'] = beta2 * state['den_sum'] + (1.0 - beta2)
+
+                de_nom = exp_avg_sq.sqrt().add_(math.sqrt(state['den_sum']) * group['eps'])
+
+                update = (group['alpha'] * grad + exp_avg).div_(de_nom).div_(math.sqrt(state['den_sum']))
+
+                p.add_(update, alpha=-group['lr'])
+
+        return loss