Add Lion optimizer (#504)

* Create lion.py

Add the Lion optimizer from https://arxiv.org/pdf/2302.06675.pdf

* Update __init__.py

* Update test_basic.py

* Update test_optimizer.py

* Update test_optimizer_with_nn.py

* Update test_optimizer_with_nn.py

* Update lion.py

* Update test_optimizer_with_nn.py

* Update test_param_validation.py

Author: jona-sassenhagen

tests/test_basic.py

@@ -70,9 +70,9 @@ def build_lookahead(*a, **kw):
     (optim.Adahessian, {'lr': 0.15, 'hessian_power': 0.6, 'seed': 0}, 900),
     (optim.MADGRAD, {'lr': 0.02}, 500),
     (optim.LARS, {'lr': 0.002, 'momentum': 0.91}, 900),
+    (optim.Lion, {'lr': 0.025}, 3600),
 ]
 
-
 @pytest.mark.parametrize('case', cases, ids=ids)
 @pytest.mark.parametrize('optimizer_config', optimizers, ids=ids)
 def test_benchmark_function(case, optimizer_config):

tests/test_optimizer.py

@@ -94,6 +94,7 @@ def build_lookahead(*a, **kw):
     optim.SWATS,
     optim.Shampoo,
     optim.Yogi,
+    optim.Lion,
 ]
 
 

tests/test_optimizer_with_nn.py

@@ -89,6 +89,7 @@ def build_lookahead(*a, **kw):
     ),
     (optim.Yogi, {'lr': 0.1, 'weight_decay': 1e-3}, 200),
     (optim.Adahessian, {'lr': 0.1, 'weight_decay': 1e-3}, 200),
+    (optim.Lion, {'lr': 0.1, 'weight_decay': 1e-3}, 200),
 ]
 
 

tests/test_param_validation.py

@@ -55,6 +55,7 @@ def test_sparse_not_supported(optimizer_class):
     optim.SWATS,
     optim.Shampoo,
     optim.Yogi,
+    optim.Lion,
 ]
 
 
@@ -118,6 +119,7 @@ def test_eps_validation(optimizer_class):
     optim.SWATS,
     optim.Shampoo,
     optim.Yogi,
+    optim.Lion,
 ]
 
 
@@ -141,6 +143,7 @@ def test_weight_decay_validation(optimizer_class):
     optim.QHAdam,
     optim.RAdam,
     optim.Yogi,
+    optim.Lion,
 ]
 
 

torch_optimizer/__init__.py

@@ -44,6 +44,7 @@
 from .shampoo import Shampoo
 from .swats import SWATS
 from .yogi import Yogi
+from .lion import Lion
 
 __all__ = (
     'A2GradExp',
@@ -76,6 +77,7 @@
     'SWATS',
     'Shampoo',
     'Yogi',
+    'Lion',
     # utils
     'get',
 )
@@ -107,6 +109,7 @@
     SWATS,
     Shampoo,
     Yogi,
+    Lion,
 ]  # type: List[Type[Optimizer]]
 
 

torch_optimizer/lion.py

@@ -0,0 +1,95 @@
+import torch
+from torch.optim.optimizer import Optimizer
+
+from .types import OptFloat, OptLossClosure, Params, Betas2
+
+__all__ = ("Lion",)
+
+
+class Lion(Optimizer):
+    r"""Implements Lion algorithm.
+
+    Addapted from https://github.com/google/automl/tree/master/lion
+
+    The Lion - EvoLved SIgn MOmeNtum - algorithm was proposed in
+    https://arxiv.org/pdf/2302.06675.pdf.
+    Lion aims to be more memory efficient than Adam by only tracking momentum.
+
+    Caveats: As detailed in the paper, Lion requires a smaller learning rate lr,
+    and larger decoupled weight decay to maintain effective weight decay strength.
+    Also, the gain of Lion increases with the batch size.
+    Furthermore, Lion was not found to outperform AdamW on some large language
+    and text/image datasets.
+
+    Arguments:
+        params: iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr: learning rate (default: 1e-3)
+        betas: coefficients used for computing
+            running averages of gradient and its square (default: (0.95, 0))
+        weight_decay: weight decay (L2 penalty) (default: 0)
+
+    Example:
+        >>> import torch_optimizer as optim
+        >>> optimizer = optim.Lion(model.parameters(), lr=0.001)
+        >>> optimizer.zero_grad()
+        >>> loss_fn(model(input), target).backward()
+        >>> optimizer.step()
+    """
+
+    def __init__(
+        self,
+        params: Params,
+        lr: float = 1e-4,
+        betas: Betas2 = (0.9, 0.99),
+        weight_decay: float = 0.0,
+    ):
+
+        if lr <= 0.0:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+        if weight_decay < 0:
+            raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
+        defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay)
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self, closure: OptLossClosure = None) -> OptFloat:
+        r"""Performs a single optimization step.
+
+        Arguments:
+            closure: A closure that reevaluates the model and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+
+                # Perform stepweight decay
+                p.data.mul_(1 - group["lr"] * group["weight_decay"])
+
+                grad = p.grad
+                state = self.state[p]
+                # State initialization
+                if len(state) == 0:
+                    # Exponential moving average of gradient values
+                    state["exp_avg"] = torch.zeros_like(p)
+
+                exp_avg = state["exp_avg"]
+                beta1, beta2 = group["betas"]
+
+                # Weight update
+                update = exp_avg * beta1 + grad * (1 - beta1)
+                p.add_(torch.sign(update), alpha=-group["lr"])
+                # Decay the momentum running average coefficient
+                exp_avg.mul_(beta2).add_(grad, alpha=1 - beta2)
+
+        return loss