Merge pull request #385 from kozistr/fix/scion-optimizer

[Fix] momentum in SCION optimizer

Author: kozistr

docs/changelogs/v3.6.1.md

@@ -0,0 +1,5 @@
+## Change Log
+
+### Fix
+
+* Fix to use `momentum buffer` instead of the gradient to calculate LMO. (#385)

pytorch_optimizer/optimizer/scion.py

@@ -409,7 +409,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
 
                 d.mul_(1.0 - group['momentum']).add_(grad, alpha=group['momentum'])
 
-                update = norm.lmo(grad).mul_(group['scale'])
+                update = norm.lmo(d).mul_(group['scale'])
 
                 if group['constraint']:
                     p.mul_(1.0 - group['lr'])

tests/constants.py

@@ -640,7 +640,7 @@
     (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}, 5),
-    (SCION, {'lr': 2e-1, 'constraint': True}, 5),
+    (SCION, {'lr': 5e-1, 'constraint': True}, 20),
     (SCIONLight, {'lr': 5e-1, 'constraint': False, 'weight_decay': 1e-3}, 5),
     (SCIONLight, {'lr': 2e-1, 'constraint': True}, 5),
     (AdaGC, {'lr': 1e-1, 'warmup_steps': 2, 'weight_decay': 1e-3}, 5),