pytorch-optimizer 3.4.1__tar.gz → 3.4.2__tar.gz

{pytorch_optimizer-3.4.1 → pytorch_optimizer-3.4.2}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.3
 Name: pytorch_optimizer
-Version: 3.4.1
+Version: 3.4.2
 Summary: optimizer & lr scheduler & objective function collections in PyTorch
 License: Apache-2.0
-Keywords: pytorch,deep-learning,optimizer,lr scheduler,A2Grad,ASGD,AccSGD,AdaBelief,AdaBound,AdaDelta,AdaFactor,AdaMax,AdamG,AdaMod,AdaNorm,AdaPNM,AdaSmooth,AdEMAMix,ADOPT,AdaHessian,Adai,Adalite,AdaLomo,AdamMini,AdamP,AdamS,Adan,AggMo,Aida,AliG,Amos,Apollo,APOLLO,AvaGrad,bSAM,CAME,DAdaptAdaGrad,DAdaptAdam,DAdaptAdan,DAdaptSGD,DAdaptLion,DeMo,DiffGrad,EXAdam,FAdam,FOCUS,Fromage,FTRL,GaLore,Grams,Gravity,GrokFast,GSAM,Kate,Lamb,LaProp,LARS,Lion,LOMO,Lookahead,MADGRAD,MARS,MSVAG,Muno,Nero,NovoGrad,OrthoGrad,PAdam,PCGrad,PID,PNM,Prodigy,PSGD,QHAdam,QHM,RAdam,Ranger,Ranger21,RotoGrad,SAM,GCSAM,LookSAM,ScheduleFreeSGD,ScheduleFreeAdamW,ScheduleFreeRAdam,SGDP,Shampoo,ScalableShampoo,SGDW,SignSGD,SM3,SOAP,SopihaH,SPAM,SRMM,StableAdamW,SWATS,TAM,Tiger,TRAC,WSAM,Yogi,BCE,BCEFocal,Focal,FocalCosine,SoftF1,Dice,LDAM,Jaccard,Bi-Tempered,Tversky,FocalTversky,LovaszHinge,bitsandbytes,WSD,QGaLore
+Keywords: pytorch,deep-learning,optimizer,lr scheduler,A2Grad,ASGD,AccSGD,AdaBelief,AdaBound,AdaDelta,AdaFactor,AdaMax,AdamG,AdaMod,AdaNorm,AdaPNM,AdaSmooth,AdEMAMix,ADOPT,AdaHessian,Adai,Adalite,AdaLomo,AdamMini,AdamP,AdamS,Adan,AggMo,Aida,AliG,Amos,Apollo,APOLLO,AvaGrad,bSAM,CAME,DAdaptAdaGrad,DAdaptAdam,DAdaptAdan,DAdaptSGD,DAdaptLion,DeMo,DiffGrad,EXAdam,FAdam,FOCUS,Fromage,FTRL,GaLore,Grams,Gravity,GrokFast,GSAM,Kate,Lamb,LaProp,LARS,Lion,LOMO,Lookahead,MADGRAD,MARS,MSVAG,Muno,Nero,NovoGrad,OrthoGrad,PAdam,PCGrad,PID,PNM,Prodigy,PSGD,QHAdam,QHM,RAdam,Ranger,Ranger21,RotoGrad,SAM,GCSAM,LookSAM,ScheduleFreeSGD,ScheduleFreeAdamW,ScheduleFreeRAdam,SCION,SGDP,Shampoo,ScalableShampoo,SGDW,SignSGD,SM3,SOAP,SopihaH,SPAM,SRMM,StableAdamW,SWATS,TAM,Tiger,TRAC,WSAM,Yogi,BCE,BCEFocal,Focal,FocalCosine,SoftF1,Dice,LDAM,Jaccard,Bi-Tempered,Tversky,FocalTversky,LovaszHinge,bitsandbytes,WSD,QGaLore
 Author: kozistr
 Author-email: kozistr@gmail.com
 Maintainer: kozistr
@@ -48,7 +48,7 @@ Description-Content-Type: text/markdown
 
 ## The reasons why you use `pytorch-optimizer`.
 
-* Wide range of supported optimizers. Currently, **98 optimizers (+ `bitsandbytes`, `qgalore`, `torchao`)**, **16 lr schedulers**, and **13 loss functions** are supported!
+* Wide range of supported optimizers. Currently, **99 optimizers (+ `bitsandbytes`, `qgalore`, `torchao`)**, **16 lr schedulers**, and **13 loss functions** are supported!
 * Including many variants such as `ADOPT`, `Cautious`, `AdamD`, `StableAdamW`, and `Gradient Centrailiaztion`
 * Easy to use, clean, and tested codes
 * Active maintenance
@@ -244,6 +244,7 @@ get_supported_optimizers(['adam*', 'ranger*'])
 | EXAdam        | *The Power of Adaptive Cross-Moments*                                                             | [github](https://github.com/AhmedMostafa16/EXAdam)                                                             | <https://arxiv.org/abs/2412.20302>                                                          | [cite](https://github.com/AhmedMostafa16/EXAdam?tab=readme-ov-file#citation)                                                        |
 | GCSAM         | *Gradient Centralized Sharpness Aware Minimization*                                               | [github](https://github.com/mhassann22/GCSAM)                                                                  | <https://arxiv.org/abs/2501.11584>                                                          | [cite](https://github.com/mhassann22/GCSAM?tab=readme-ov-file#citation)                                                             |
 | LookSAM       | *Towards Efficient and Scalable Sharpness-Aware Minimization*                                     | [github](https://github.com/rollovd/LookSAM)                                                                   | <https://arxiv.org/abs/2203.02714>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2022arXiv220302714L/exportcitation)                                                        |
+| SCION         | *Training Deep Learning Models with Norm-Constrained LMOs*                                        |                                                                                                                | <https://arxiv.org/abs/2502.07529>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2025arXiv250207529P/exportcitation)                                                        |
 
 ## Supported LR Scheduler
 

{pytorch_optimizer-3.4.1 → pytorch_optimizer-3.4.2}/README.md

@@ -10,7 +10,7 @@
 
 ## The reasons why you use `pytorch-optimizer`.
 
-* Wide range of supported optimizers. Currently, **98 optimizers (+ `bitsandbytes`, `qgalore`, `torchao`)**, **16 lr schedulers**, and **13 loss functions** are supported!
+* Wide range of supported optimizers. Currently, **99 optimizers (+ `bitsandbytes`, `qgalore`, `torchao`)**, **16 lr schedulers**, and **13 loss functions** are supported!
 * Including many variants such as `ADOPT`, `Cautious`, `AdamD`, `StableAdamW`, and `Gradient Centrailiaztion`
 * Easy to use, clean, and tested codes
 * Active maintenance
@@ -206,6 +206,7 @@ get_supported_optimizers(['adam*', 'ranger*'])
 | EXAdam        | *The Power of Adaptive Cross-Moments*                                                             | [github](https://github.com/AhmedMostafa16/EXAdam)                                                             | <https://arxiv.org/abs/2412.20302>                                                          | [cite](https://github.com/AhmedMostafa16/EXAdam?tab=readme-ov-file#citation)                                                        |
 | GCSAM         | *Gradient Centralized Sharpness Aware Minimization*                                               | [github](https://github.com/mhassann22/GCSAM)                                                                  | <https://arxiv.org/abs/2501.11584>                                                          | [cite](https://github.com/mhassann22/GCSAM?tab=readme-ov-file#citation)                                                             |
 | LookSAM       | *Towards Efficient and Scalable Sharpness-Aware Minimization*                                     | [github](https://github.com/rollovd/LookSAM)                                                                   | <https://arxiv.org/abs/2203.02714>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2022arXiv220302714L/exportcitation)                                                        |
+| SCION         | *Training Deep Learning Models with Norm-Constrained LMOs*                                        |                                                                                                                | <https://arxiv.org/abs/2502.07529>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2025arXiv250207529P/exportcitation)                                                        |
 
 ## Supported LR Scheduler
 

{pytorch_optimizer-3.4.1 → pytorch_optimizer-3.4.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytorch_optimizer"
-version = "3.4.1"
+version = "3.4.2"
 description = "optimizer & lr scheduler & objective function collections in PyTorch"
 license = "Apache-2.0"
 authors = ["kozistr <kozistr@gmail.com>"]
@@ -18,10 +18,10 @@ keywords = [
     "GrokFast", "GSAM", "Kate", "Lamb", "LaProp", "LARS", "Lion", "LOMO", "Lookahead", "MADGRAD", "MARS", "MSVAG",
     "Muno", "Nero", "NovoGrad", "OrthoGrad", "PAdam", "PCGrad", "PID", "PNM", "Prodigy", "PSGD", "QHAdam", "QHM",
     "RAdam", "Ranger", "Ranger21", "RotoGrad", "SAM", "GCSAM", "LookSAM", "ScheduleFreeSGD", "ScheduleFreeAdamW",
-    "ScheduleFreeRAdam", "SGDP", "Shampoo", "ScalableShampoo", "SGDW", "SignSGD", "SM3", "SOAP", "SopihaH", "SPAM",
-    "SRMM", "StableAdamW", "SWATS", "TAM", "Tiger", "TRAC", "WSAM", "Yogi", "BCE", "BCEFocal", "Focal", "FocalCosine",
-    "SoftF1", "Dice", "LDAM", "Jaccard", "Bi-Tempered", "Tversky", "FocalTversky", "LovaszHinge", "bitsandbytes", "WSD",
-    "QGaLore",
+    "ScheduleFreeRAdam", "SCION", "SGDP", "Shampoo", "ScalableShampoo", "SGDW", "SignSGD", "SM3", "SOAP", "SopihaH",
+    "SPAM", "SRMM", "StableAdamW", "SWATS", "TAM", "Tiger", "TRAC", "WSAM", "Yogi", "BCE", "BCEFocal", "Focal",
+    "FocalCosine", "SoftF1", "Dice", "LDAM", "Jaccard", "Bi-Tempered", "Tversky", "FocalTversky", "LovaszHinge",
+    "bitsandbytes", "WSD", "QGaLore",
 ]
 classifiers = [
     "License :: OSI Approved :: Apache Software License",

{pytorch_optimizer-3.4.1 → pytorch_optimizer-3.4.2}/pytorch_optimizer/__init__.py

@@ -58,6 +58,7 @@ from pytorch_optimizer.optimizer import (
     PNM,
     QHM,
     SAM,
+    SCION,
     SGDP,
     SGDW,
     SM3,

{pytorch_optimizer-3.4.1 → pytorch_optimizer-3.4.2}/pytorch_optimizer/optimizer/__init__.py

@@ -81,6 +81,7 @@ from pytorch_optimizer.optimizer.ranger21 import Ranger21
 from pytorch_optimizer.optimizer.rotograd import RotoGrad
 from pytorch_optimizer.optimizer.sam import BSAM, GSAM, SAM, WSAM, LookSAM
 from pytorch_optimizer.optimizer.schedulefree import ScheduleFreeAdamW, ScheduleFreeRAdam, ScheduleFreeSGD
+from pytorch_optimizer.optimizer.scion import SCION
 from pytorch_optimizer.optimizer.sgd import ASGD, SGDW, AccSGD, SGDSaI, SignSGD
 from pytorch_optimizer.optimizer.sgdp import SGDP
 from pytorch_optimizer.optimizer.shampoo import ScalableShampoo, Shampoo
@@ -300,6 +301,7 @@ OPTIMIZER_LIST: List[OPTIMIZER] = [
     SPAM,
     Kron,
     EXAdam,
+    SCION,
     Ranger25,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}

{pytorch_optimizer-3.4.1 → pytorch_optimizer-3.4.2}/pytorch_optimizer/optimizer/experimental/ranger25.py

@@ -40,7 +40,7 @@ class Ranger25(BaseOptimizer):
         self,
         params: PARAMETERS,
         lr: float = 1e-3,
-        betas: BETAS = (0.95, 0.98, 0.9999),
+        betas: BETAS = (0.9, 0.98, 0.9999),
         weight_decay: float = 1e-3,
         weight_decouple: bool = True,
         fixed_decay: bool = False,

{pytorch_optimizer-3.4.1 → pytorch_optimizer-3.4.2}/pytorch_optimizer/optimizer/schedulefree.py

@@ -1,4 +1,3 @@
-import math
 from typing import List
 
 import torch
@@ -15,8 +14,6 @@ class ScheduleFreeSGD(BaseOptimizer):
     :param lr: float. learning rate.
     :param momentum: float. momentum factor, must be between 0 and 1 exclusive.
     :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 r: float. use polynomial weighting in the average with power r.
     :param weight_lr_power: float. during warmup, the weights in the average will be equal to lr raised to this power.
         set to 0 for no weighting.
@@ -30,8 +27,6 @@ class ScheduleFreeSGD(BaseOptimizer):
         lr: float = 1.0,
         momentum: float = 0.9,
         weight_decay: float = 0.0,
-        weight_decouple: bool = True,
-        fixed_decay: bool = False,
         r: float = 0.0,
         weight_lr_power: float = 2.0,
         warmup_steps: int = 0,
@@ -47,8 +42,6 @@ class ScheduleFreeSGD(BaseOptimizer):
             'lr': lr,
             'momentum': momentum,
             'weight_decay': weight_decay,
-            'weight_decouple': weight_decouple,
-            'fixed_decay': fixed_decay,
             'r': r,
             'weight_lr_power': weight_lr_power,
             'warmup_steps': warmup_steps,
@@ -114,7 +107,7 @@ class ScheduleFreeSGD(BaseOptimizer):
             lr: float = group['lr'] * schedule
             lr_max = group['lr_max'] = max(lr, group['lr_max'])
 
-            weight = (group['step'] ** group['r']) * (lr_max ** group['weight_lr_power'])
+            weight: float = (group['step'] ** group['r']) * (lr_max ** group['weight_lr_power'])
             weight_sum = group['weight_sum'] = group['weight_sum'] + weight
 
             checkpoint: float = weight / weight_sum if weight_sum != 0.0 else 0.0
@@ -137,8 +130,8 @@ class ScheduleFreeSGD(BaseOptimizer):
                     grad=grad,
                     lr=lr,
                     weight_decay=group['weight_decay'],
-                    weight_decouple=group['weight_decouple'],
-                    fixed_decay=group['fixed_decay'],
+                    weight_decouple=False,
+                    fixed_decay=False,
                 )
 
                 z = state['z']
@@ -158,8 +151,6 @@ class ScheduleFreeAdamW(BaseOptimizer):
     :param lr: float. learning rate.
     :param betas: BETAS. coefficients used for computing running averages of gradient and the squared hessian trace.
     :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 r: float. use polynomial weighting in the average with power r.
     :param weight_lr_power: float. during warmup, the weights in the average will be equal to lr raised to this power.
         set to 0 for no weighting.
@@ -174,8 +165,6 @@ class ScheduleFreeAdamW(BaseOptimizer):
         lr: float = 2.5e-3,
         betas: BETAS = (0.9, 0.999),
         weight_decay: float = 0.0,
-        weight_decouple: bool = True,
-        fixed_decay: bool = False,
         r: float = 0.0,
         weight_lr_power: float = 2.0,
         warmup_steps: int = 0,
@@ -192,8 +181,6 @@ class ScheduleFreeAdamW(BaseOptimizer):
             'lr': lr,
             'betas': betas,
             'weight_decay': weight_decay,
-            'weight_decouple': weight_decouple,
-            'fixed_decay': fixed_decay,
             'r': r,
             'weight_lr_power': weight_lr_power,
             'warmup_steps': warmup_steps,
@@ -259,22 +246,16 @@ class ScheduleFreeAdamW(BaseOptimizer):
 
             beta1, beta2 = group['betas']
 
-            bias_correction2_sq: float = math.sqrt(1.0 - beta2 ** group['step'])
+            bias_correction2: float = self.debias(beta2, group['step'])
 
-            lr: float = group['lr'] * schedule * bias_correction2_sq
+            lr: float = group['lr'] * schedule
             lr_max = group['lr_max'] = max(lr, group['lr_max'])
 
-            weight = (group['step'] ** group['r']) * (lr_max ** group['weight_lr_power'])
+            weight: float = (group['step'] ** group['r']) * (lr_max ** group['weight_lr_power'])
             weight_sum = group['weight_sum'] = group['weight_sum'] + weight
 
             checkpoint: float = weight / weight_sum if weight_sum != 0.0 else 0.0
 
-            if group['use_palm']:
-                beta2: float = 1.0 - group['step'] ** -0.8
-                debias: float = (1.0 - beta2) / (1.0 - beta2 ** group['step'])
-            else:
-                debias: float = beta2
-
             for p in group['params']:
                 if p.grad is None:
                     continue
@@ -289,27 +270,27 @@ class ScheduleFreeAdamW(BaseOptimizer):
                     state['z'] = p.clone()
                     state['exp_avg_sq'] = torch.zeros_like(p)
 
-                self.apply_weight_decay(
-                    p=p,
-                    grad=grad,
-                    lr=lr,
-                    weight_decay=group['weight_decay'],
-                    weight_decouple=group['weight_decouple'],
-                    fixed_decay=group['fixed_decay'],
-                )
-
                 z, exp_avg_sq = state['z'], state['exp_avg_sq']
-                exp_avg_sq.mul_(debias).addcmul_(grad, grad, value=1.0 - debias)
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
 
                 de_nom = self.apply_ams_bound(
                     ams_bound=group['ams_bound'],
-                    exp_avg_sq=exp_avg_sq,
+                    exp_avg_sq=exp_avg_sq.div(bias_correction2),
                     max_exp_avg_sq=state.get('max_exp_avg_sq', None),
                     eps=group['eps'],
                 )
 
                 grad.div_(de_nom)
 
+                self.apply_weight_decay(
+                    p=p,
+                    grad=grad,
+                    lr=lr,
+                    weight_decay=group['weight_decay'],
+                    weight_decouple=False,
+                    fixed_decay=False,
+                )
+
                 p.lerp_(z, weight=checkpoint)
                 p.add_(grad, alpha=lr * (beta1 * (1.0 - checkpoint) - 1))
 
@@ -325,12 +306,13 @@ class ScheduleFreeRAdam(BaseOptimizer):
     :param lr: float. learning rate.
     :param betas: BETAS. coefficients used for computing running averages of gradient and the squared hessian trace.
     :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 degenerated_to_sgd: float. degenerated to SGD.
     :param r: float. use polynomial weighting in the average with power r.
     :param weight_lr_power: float. during warmup, the weights in the average will be equal to lr raised to this power.
         set to 0 for no weighting.
+    :param silent_sgd_phase: bool. the optimizer will not use the first SGD phase of RAdam. This means that the
+        optimizer will not update model parameters during the early training steps (e.g., < 5 when β_2 = 0.999), but
+        just update the momentum values of the optimizer. This helps stabilize training by ensuring smoother warmup
+        behavior and more reliable calculation of the moving average coefficient (`ckp1`). Recommended to set to True.
     :param eps: float. term added to the denominator to improve numerical stability.
     """
 
@@ -340,11 +322,9 @@ class ScheduleFreeRAdam(BaseOptimizer):
         lr: float = 2.5e-3,
         betas: BETAS = (0.9, 0.999),
         weight_decay: float = 0.0,
-        weight_decouple: bool = True,
-        fixed_decay: bool = False,
-        degenerated_to_sgd: bool = False,
         r: float = 0.0,
         weight_lr_power: float = 2.0,
+        silent_sgd_phase: bool = True,
         eps: float = 1e-8,
         **kwargs,
     ):
@@ -357,9 +337,7 @@ class ScheduleFreeRAdam(BaseOptimizer):
             'lr': lr,
             'betas': betas,
             'weight_decay': weight_decay,
-            'weight_decouple': weight_decouple,
-            'fixed_decay': fixed_decay,
-            'degenerated_to_sgd': degenerated_to_sgd,
+            'silent_sgd_phase': silent_sgd_phase,
             'r': r,
             'weight_lr_power': weight_lr_power,
             'eps': eps,
@@ -418,7 +396,7 @@ class ScheduleFreeRAdam(BaseOptimizer):
 
             beta1, beta2 = group['betas']
 
-            bias_correction2_sq: float = math.sqrt(1.0 - beta2 ** group['step'])
+            bias_correction2: float = self.debias_beta(beta2, group['step'])
 
             lr, n_sma = self.get_rectify_step_size(
                 is_rectify=True,
@@ -426,24 +404,20 @@ class ScheduleFreeRAdam(BaseOptimizer):
                 lr=group['lr'],
                 beta2=beta2,
                 n_sma_threshold=4,
-                degenerated_to_sgd=group['degenerated_to_sgd'],
+                degenerated_to_sgd=False,
             )
+            if lr < 0.0:
+                lr = float(not group['silent_sgd_phase'])
 
             lr_max = group['lr_max'] = max(lr, group['lr_max'])
 
-            weight = (group['step'] ** group['r']) * (lr_max ** group['weight_lr_power'])
+            weight: float = (group['step'] ** group['r']) * (lr_max ** group['weight_lr_power'])
             weight_sum = group['weight_sum'] = group['weight_sum'] + weight
 
             checkpoint: float = weight / weight_sum if weight_sum != 0.0 else 0.0
 
             adaptive_y_lr: float = lr * (beta1 * (1.0 - checkpoint) - 1.0)
 
-            if group['use_palm']:
-                beta2: float = 1.0 - group['step'] ** -0.8
-                debias: float = (1.0 - beta2) / (1.0 - beta2 ** group['step'])
-            else:
-                debias: float = beta2
-
             for p in group['params']:
                 if p.grad is None:
                     continue
@@ -459,10 +433,10 @@ class ScheduleFreeRAdam(BaseOptimizer):
                     state['exp_avg_sq'] = torch.zeros_like(p)
 
                 z, exp_avg_sq = state['z'], state['exp_avg_sq']
-                exp_avg_sq.mul_(debias).addcmul_(grad, grad, value=1.0 - debias)
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
 
                 if n_sma > 4.0:
-                    de_nom = exp_avg_sq.sqrt().div_(bias_correction2_sq).add_(group['eps'])
+                    de_nom = exp_avg_sq.sqrt().div_(bias_correction2).add_(group['eps'])
                     grad.div_(de_nom)
 
                 self.apply_weight_decay(
@@ -470,8 +444,8 @@ class ScheduleFreeRAdam(BaseOptimizer):
                     grad=grad,
                     lr=lr,
                     weight_decay=group['weight_decay'],
-                    weight_decouple=group['weight_decouple'],
-                    fixed_decay=group['fixed_decay'],
+                    weight_decouple=False,
+                    fixed_decay=False,
                 )
 
                 p.lerp_(z, weight=checkpoint)

pytorch_optimizer-3.4.2/pytorch_optimizer/optimizer/scion.py

@@ -0,0 +1,113 @@
+from typing import Literal
+
+import torch
+
+from pytorch_optimizer.base.exception import NoSparseGradientError
+from pytorch_optimizer.base.optimizer import BaseOptimizer
+from pytorch_optimizer.base.type import CLOSURE, DEFAULTS, LOSS, PARAMETERS
+from pytorch_optimizer.optimizer.shampoo_utils import zero_power_via_newton_schulz_5
+
+LMO_TYPE = Literal['spectral', 'sign', 'col_norm', 'row_norm']
+
+
+class SCION(BaseOptimizer):
+    r"""Training Deep Learning Models with Norm-Constrained LMOs.
+
+    :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
+    :param lr: float. learning rate.
+    :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 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,
+        momentum: float = 0.1,
+        constraint: bool = False,
+        lmo_type: LMO_TYPE = 'spectral',
+        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_options(lmo_type, 'lmo_type', ['spectral', 'sign', 'col_norm', 'row_norm'])
+
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'momentum': momentum,
+            'constraint': constraint,
+            'lmo_type': lmo_type,
+            'weight_decay': weight_decay,
+            'weight_decouple': weight_decouple,
+        }
+        super().__init__(params, defaults)
+
+    def __str__(self) -> str:
+        return 'SCION'
+
+    @torch.no_grad()
+    def reset(self):
+        for group in self.param_groups:
+            for p in group['params']:
+                state = self.state[p]
+                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)
+        if lmo_type == 'sign':
+            return torch.sign(grad)
+        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)
+
+    @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:
+            step_size: float = -group['lr']
+            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 'd' not in state:
+                    state['d'] = torch.zeros_like(p)
+
+                d = state['d']
+                d.mul_(1.0 - group['momentum']).add_(grad, alpha=group['momentum'])
+
+                update = self.get_lmo_direction(d, group['lmo_type'])
+
+                if not group['constraint']:
+                    self.apply_weight_decay(
+                        p,
+                        grad,
+                        lr=group['lr'],
+                        weight_decay=group['weight_decay'],
+                        weight_decouple=group['weight_decouple'],
+                        fixed_decay=False,
+                    )
+
+                    p.add_(update, alpha=step_size)
+                else:
+                    p.mul_(1.0 - step_size).add_(update, alpha=step_size)
+
+        return loss