Add sparse lion and use as default

config/high_order_interpolation.yaml

@@ -67,7 +67,7 @@ prompts:
 topk: 3
 num_predict: 100
 defaults:
-  - optimizer: lion
+  - optimizer: sparse_lion
   - net: small
   #- override hydra/sweeper: nevergrad
 # hydra:

config/optimizer/sparse_lion.yaml

@@ -0,0 +1,7 @@
+name: sparse_lion
+# lr supposed to be 3 to 5 times smaller than adam
+lr: 1e-5
+patience: 5
+factor: 0.1
+gamma: 0.9
+scheduler: plateau # exponential

language_interpolation/networks.py

@@ -19,6 +19,7 @@
 import logging
 import time
 from lion_pytorch import Lion
+from high_order_layers_torch.sparse_optimizers import SparseLion
 
 logger = logging.getLogger(__name__)
 logging.basicConfig(level=logging.DEBUG)
@@ -160,6 +161,10 @@ def configure_optimizers(self):
             optimizer = Lion(
                 self.parameters(), lr=self.cfg.optimizer.lr, weight_decay=0.0
             )
+        elif self.cfg.optimizer.name == "sparse_lion":
+            optimizer = SparseLion(
+                self.parameters(), lr=self.cfg.optimizer.lr, weight_decay=0.0
+            )
         elif self.cfg.optimizer.name == "adam":
             optimizer = optim.Adam(
                 params=self.parameters(),
@@ -798,7 +803,7 @@ def select_network(cfg: DictConfig, device: str = None):
             hidden_segments=cfg.net.hidden.segments,
             normalization=normalization,
             device=cfg.accelerator,
-            layer_type_in=cfg.net.input.get('layer_type', None),
+            layer_type_in=cfg.net.input.get("layer_type", None),
         )
     elif cfg.net.model_type == "high_order_conv":
         conv = HighOrderFullyConvolutionalNetwork(

pyproject.toml

@@ -18,7 +18,7 @@ langchain = "^0.0.334"
 hydra-core = "^1.3.2"
 tensorboard = "^2.15.1"
 lion-pytorch = "^0.1.2"
-high-order-layers-torch = "^2.3.3"
+high-order-layers-torch = "^2.4.0"
 
 [tool.poetry.group.dev.dependencies]
 black = "^23.11.0"