Merge pull request #362 from kozistr/update/scion-optimizer

[Update] SCION optimizer

Author: kozistr

docs/changelogs/v3.4.3.md

@@ -13,6 +13,9 @@
     * adjust default hyperparameters the same as the original implementation.
     * support adjusted lr from the Moonlight. you can use it by setting `use_adjusted_lr=True`.
 * Tune the performance of the coupled Newton iteration method by 5% increase. (#360)
+* Update `SCION` optimizer. (#361)
+    * add `scale` parameter.
+    * update `get_lmo_direction`.
 
 ### Fix
 

docs/visualization.md

@@ -350,6 +350,10 @@
 
 ![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rastrigin_StableAdamW.png)
 
+### StableSPAM
+
+![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rastrigin_StableSPAM.png)
+
 ### SWATS
 
 ![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rastrigin_SWATS.png)
@@ -716,6 +720,10 @@
 
 ![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rosenbrock_StableAdamW.png)
 
+### StableSPAM
+
+![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rosenbrock_StableSPAM.png)
+
 ### SWATS
 
 ![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rosenbrock_SWATS.png)

docs/visualizations/rastrigin_SCION.png

@@ -204,7 +204,6 @@ def closure() -> float:
     parameters = list(model.parameters())
     optimizer_name: str = optimizer_class.__name__.lower()
 
-    # Special handling for optimizers with unique requirements
     if optimizer_name == 'ranger21':
         optimizer_config['num_iterations'] = num_iters
     elif optimizer_name == 'ranger25':
@@ -215,11 +214,12 @@ def closure() -> float:
         optimizer_config['projection_fn'] = lambda: l2_projection(parameters, max_norm=1)
     elif optimizer_name == 'bsam':
         optimizer_config['num_data'] = 1
+    elif optimizer_name == 'scion':
+        optimizer_config['scale'] = 50.0
 
-    if optimizer_name in OPTIMIZERS_MODEL_INPUT_NEEDED:
-        optimizer = optimizer_class(model, **optimizer_config)
-    else:
-        optimizer = optimizer_class(parameters, **optimizer_config)
+    optimizer = optimizer_class(
+        model if optimizer_name in OPTIMIZERS_MODEL_INPUT_NEEDED else parameters, **optimizer_config
+    )
 
     steps = torch.zeros((2, num_iters + 1), dtype=torch.float32)
     steps[:, 0] = model.x.detach()
@@ -394,7 +394,7 @@ def execute_experiments(
                 rstate=np.random.default_rng(seed),
             )
         except AllTrialsFailed:
-            print(f'⚠️ {optimizer_name} failed to optimize {func.__name__}')  # noqa: T201
+            print(f'{optimizer_name} failed to optimize {func.__name__}')  # noqa: T201
             continue
 
         steps, _ = execute_steps(func, initial_state, optimizer_class, best_params.copy(), TESTING_OPTIMIZATION_STEPS)

docs/visualizations/rastrigin_StableSPAM.png

@@ -1,3 +1,4 @@
+import math
 from typing import Literal
 
 import torch
@@ -18,30 +19,35 @@ class SCION(BaseOptimizer):
     :param momentum: float. momentum factor.
     :param constraint: bool. whether to use a constraint SCG or not.
     :param lmo_type: LMO_TYPE. supported LMO types.
+    :param scale: float. based on the usage of the original intend, 50.0 is used for Transformer block, and 3000.0 is
+        used for others (e.g. Embedding, LM head)
     :param weight_decay: float. weight decay (L2 penalty).
     :param weight_decouple: bool. the optimizer uses decoupled weight decay as in AdamW.
     """
 
     def __init__(
         self,
         params: PARAMETERS,
-        lr: float = 1e-4,
+        lr: float = 1e-3,
         momentum: float = 0.1,
         constraint: bool = False,
         lmo_type: LMO_TYPE = 'spectral',
+        scale: float = 1.0,
         weight_decay: float = 0.0,
         weight_decouple: bool = True,
         **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_range(momentum, 'momentum', 0.0, 1.0, '(]')
+        self.validate_positive(scale, 'scale')
         self.validate_options(lmo_type, 'lmo_type', ['spectral', 'sign', 'col_norm', 'row_norm'])
 
         defaults: DEFAULTS = {
             'lr': lr,
             'momentum': momentum,
             'constraint': constraint,
             'lmo_type': lmo_type,
+            'scale': scale,
             'weight_decay': weight_decay,
             'weight_decouple': weight_decouple,
         }
@@ -58,17 +64,26 @@ def reset(self):
                 state['d'] = torch.zeros_like(p)
 
     @staticmethod
-    def get_lmo_direction(grad: torch.Tensor, lmo_type: str) -> torch.Tensor:
-        r"""Get LMO direction."""
-        if lmo_type == 'spectral' and grad.ndim == 2:
-            return zero_power_via_newton_schulz_5(grad)
+    def get_lmo_direction(grad: torch.Tensor, lmo_type: LMO_TYPE) -> torch.Tensor:
+        r"""Get LMO direction.
+
+        fallback to `sign`
+        """
+        d_out, d_in, *_ = grad.shape if grad.ndim > 1 else (grad.size(0), grad.size(0))
+
+        if lmo_type == 'spectral':
+            return (
+                zero_power_via_newton_schulz_5(grad.reshape(len(grad), -1))
+                .view(grad.shape)
+                .mul_(max(1.0, math.sqrt(d_out / d_in)))
+            )
         if lmo_type == 'sign':
-            return torch.sign(grad)
+            return torch.sign(grad).div_(d_in)
         if lmo_type == 'col_norm':
             return grad / torch.norm(grad, dim=0, keepdim=True).add_(1e-6)
         if lmo_type == 'row_norm' and grad.ndim == 2:
             return grad / torch.norm(grad, dim=1, keepdim=True).add_(1e-6)
-        return torch.sign(grad)
+        return torch.sign(grad).div_(d_in)
 
     @torch.no_grad()
     def step(self, closure: CLOSURE = None) -> LOSS:
@@ -89,12 +104,13 @@ def step(self, closure: CLOSURE = None) -> LOSS:
 
                 state = self.state[p]
                 if 'd' not in state:
-                    state['d'] = torch.zeros_like(p)
+                    state['d'] = torch.zeros_like(grad)
 
                 d = state['d']
                 d.mul_(1.0 - group['momentum']).add_(grad, alpha=group['momentum'])
 
                 update = self.get_lmo_direction(d, group['lmo_type'])
+                update.mul_(group['scale'])
 
                 if not group['constraint']:
                     self.apply_weight_decay(

docs/visualizations/rosenbrock_SCION.png

@@ -565,8 +565,8 @@
     (FOCUS, {'lr': 1e-1, 'weight_decay': 1e-3}, 5),
     (Kron, {'lr': 1e0, 'weight_decay': 1e-3}, 3),
     (EXAdam, {'lr': 1e-1, 'weight_decay': 1e-3}, 5),
-    (SCION, {'lr': 5e-1, 'constraint': False, 'weight_decay': 1e-3}, 10),
-    (SCION, {'lr': 1e-1, 'constraint': True}, 10),
+    (SCION, {'lr': 5e-1, 'constraint': False, 'weight_decay': 1e-3}, 5),
+    (SCION, {'lr': 1e-1, 'constraint': True, 'lmo_type': 'col_norm'}, 10),
     (Ranger25, {'lr': 1e-1}, 3),
     (Ranger25, {'lr': 1e-1, 't_alpha_beta3': 5}, 3),
     (Ranger25, {'lr': 5e-2, 'stable_adamw': False, 'orthograd': False, 'eps': None, 'lookahead_merge_time': 2}, 3),

docs/visualizations/rosenbrock_StableSPAM.png

@@ -981,6 +981,31 @@ def test_kron_optimizer():
     optimizer.step()
 
 
+def test_scion_lmo_types():
+    grad = torch.ones(2, 2)
+
+    expected = torch.FloatTensor([[0.3438, 0.3438], [0.3438, 0.3438]]).bfloat16()
+    actual = load_optimizer('scion').get_lmo_direction(grad, 'spectral')
+
+    torch.testing.assert_close(expected, actual, rtol=1e-5, atol=1e-5)
+
+    expected = torch.FloatTensor([[0.5, 0.5], [0.5, 0.5]])
+    actual = load_optimizer('scion').get_lmo_direction(grad, 'sign')
+    torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5)
+
+    expected = torch.FloatTensor([[0.7071, 0.7071], [0.7071, 0.7071]])
+    actual = load_optimizer('scion').get_lmo_direction(grad, 'row_norm')
+    torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5)
+
+    expected = torch.FloatTensor([[0.7071, 0.7071], [0.7071, 0.7071]])
+    actual = load_optimizer('scion').get_lmo_direction(grad, 'col_norm')
+    torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5)
+
+    expected = torch.FloatTensor([[0.5, 0.5], [0.5, 0.5]])
+    actual = load_optimizer('scion').get_lmo_direction(grad, 'asdf')
+    torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5)
+
+
 def test_schedulefree_wrapper():
     model = Example()