pytorch-optimizer 2.9.1__tar.gz → 2.10.0__tar.gz

{pytorch_optimizer-2.9.1 → pytorch_optimizer-2.10.0}/PKG-INFO

@@ -1,10 +1,10 @@
 Metadata-Version: 2.1
 Name: pytorch-optimizer
-Version: 2.9.1
+Version: 2.10.0
 Summary: optimizer & lr scheduler collections in PyTorch
 Home-page: https://github.com/kozistr/pytorch_optimizer
 License: Apache-2.0
-Keywords: pytorch,deep-learning,optimizer,lr scheduler,A2Grad,ASGD,AccSGD,AdaBelief,AdaBound,AdaDelta,AdaFactor,AdaMax,AdaMod,AdaNorm,AdaPNM,AdaSmooth,Adai,AdamP,AdamS,Adan,AggMo,AliG,Apollo,AvaGrad,DAdaptAdaGrad,DAdaptAdam,DAdaptAdan,DAdaptSGD,DiffGrad,Fromage,Gravity,LARS,Lamb,Lion,MADGRAD,MSVAG,Nero,NovoGrad,PID,PNM,QHAdam,QHM,RAdam,Ranger,Ranger21,SGDP,SGDW,SM3,SRMM,SWATS,ScalableShampoo,Shampoo,Yogi,SAM,GSAM,PCGrad,RotoGrad
+Keywords: pytorch,deep-learning,optimizer,lr scheduler,A2Grad,ASGD,AccSGD,AdaBelief,AdaBound,AdaDelta,AdaFactor,AdaMax,AdaMod,AdaNorm,AdaPNM,AdaSmooth,AdaHessian,Adai,AdamP,AdamS,Adan,AggMo,AliG,Amos,Apollo,AvaGrad,DAdaptAdaGrad,DAdaptAdam,DAdaptAdan,DAdaptSGD,DiffGrad,Fromage,Gravity,LARS,Lamb,Lion,MADGRAD,MSVAG,Nero,NovoGrad,PID,PNM,QHAdam,QHM,RAdam,Ranger,Ranger21,SGDP,SGDW,SignSGD,SM3,SopihaH,SRMM,SWATS,ScalableShampoo,Shampoo,Yogi,SAM,GSAM,PCGrad,RotoGrad
 Author: kozistr
 Author-email: kozistr@gmail.com
 Maintainer: kozistr
@@ -21,13 +21,8 @@ Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.6
 Classifier: Programming Language :: Python :: 3.7
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Topic :: Scientific/Engineering
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development
@@ -59,7 +54,7 @@ pytorch-optimizer
 
 | **pytorch-optimizer** is optimizer & lr scheduler collections in PyTorch.
 | I just re-implemented (speed & memory tweaks, plug-ins) the algorithm while based on the original paper. Also, It includes useful and practical optimization ideas.
-| Currently, 50 optimizers, 6 lr schedulers are supported!
+| Currently, 54 optimizers, 6 lr schedulers are supported!
 |
 | Highly inspired by `pytorch-optimizer <https://github.com/jettify/pytorch-optimizer>`__.
 
@@ -256,6 +251,14 @@ You can check the supported optimizers & lr schedulers.
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
 | AdaDelta     | *An Adaptive Learning Rate Method*                                                                |                                                                                   | `https://arxiv.org/abs/1212.5701v1 <https://arxiv.org/abs/1212.5701v1>`__                     | `cite <https://ui.adsabs.harvard.edu/abs/2012arXiv1212.5701Z/exportcitation>`__                                      |
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
+| Amos         | *An Adam-style Optimizer with Adaptive Weight Decay towards Model-Oriented Scale*                 | `github <https://github.com/google-research/jestimator>`__                        | `https://arxiv.org/abs/2210.11693 <https://arxiv.org/abs/2210.11693>`__                       | `cite <https://ui.adsabs.harvard.edu/abs/2022arXiv221011693T/exportcitation>`__                                      |
++--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
+| SignSGD      | *Compressed Optimisation for Non-Convex Problems*                                                 | `github <https://github.com/jxbz/signSGD>`__                                      | `https://arxiv.org/abs/1802.04434 <https://arxiv.org/abs/1802.04434>`__                       | `cite <https://ui.adsabs.harvard.edu/abs/2018arXiv180204434B/exportcitation>`__                                      |
++--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
+| AdaHessian   | *An Adaptive Second Order Optimizer for Machine Learning*                                         | `github <https://github.com/amirgholami/adahessian>`__                            | `https://arxiv.org/abs/2006.00719 <https://arxiv.org/abs/2006.00719>`__                       | `cite <https://github.com/amirgholami/adahessian#citation>`__                                                        |
++--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
+| Sophia       | *A Scalable Stochastic Second-order Optimizer for Language Model Pre-training*                    | `github <https://github.com/Liuhong99/Sophia>`__                                  | `https://arxiv.org/abs/2305.14342 <https://arxiv.org/abs/2305.14342>`__                       | `cite <https://github.com/Liuhong99/Sophia>`__                                                                       |
++--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
 
 Useful Resources
 ----------------

{pytorch_optimizer-2.9.1 → pytorch_optimizer-2.10.0}/README.rst

@@ -16,7 +16,7 @@ pytorch-optimizer
 
 | **pytorch-optimizer** is optimizer & lr scheduler collections in PyTorch.
 | I just re-implemented (speed & memory tweaks, plug-ins) the algorithm while based on the original paper. Also, It includes useful and practical optimization ideas.
-| Currently, 50 optimizers, 6 lr schedulers are supported!
+| Currently, 54 optimizers, 6 lr schedulers are supported!
 |
 | Highly inspired by `pytorch-optimizer <https://github.com/jettify/pytorch-optimizer>`__.
 
@@ -213,6 +213,14 @@ You can check the supported optimizers & lr schedulers.
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
 | AdaDelta     | *An Adaptive Learning Rate Method*                                                                |                                                                                   | `https://arxiv.org/abs/1212.5701v1 <https://arxiv.org/abs/1212.5701v1>`__                     | `cite <https://ui.adsabs.harvard.edu/abs/2012arXiv1212.5701Z/exportcitation>`__                                      |
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
+| Amos         | *An Adam-style Optimizer with Adaptive Weight Decay towards Model-Oriented Scale*                 | `github <https://github.com/google-research/jestimator>`__                        | `https://arxiv.org/abs/2210.11693 <https://arxiv.org/abs/2210.11693>`__                       | `cite <https://ui.adsabs.harvard.edu/abs/2022arXiv221011693T/exportcitation>`__                                      |
++--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
+| SignSGD      | *Compressed Optimisation for Non-Convex Problems*                                                 | `github <https://github.com/jxbz/signSGD>`__                                      | `https://arxiv.org/abs/1802.04434 <https://arxiv.org/abs/1802.04434>`__                       | `cite <https://ui.adsabs.harvard.edu/abs/2018arXiv180204434B/exportcitation>`__                                      |
++--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
+| AdaHessian   | *An Adaptive Second Order Optimizer for Machine Learning*                                         | `github <https://github.com/amirgholami/adahessian>`__                            | `https://arxiv.org/abs/2006.00719 <https://arxiv.org/abs/2006.00719>`__                       | `cite <https://github.com/amirgholami/adahessian#citation>`__                                                        |
++--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
+| Sophia       | *A Scalable Stochastic Second-order Optimizer for Language Model Pre-training*                    | `github <https://github.com/Liuhong99/Sophia>`__                                  | `https://arxiv.org/abs/2305.14342 <https://arxiv.org/abs/2305.14342>`__                       | `cite <https://github.com/Liuhong99/Sophia>`__                                                                       |
++--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
 
 Useful Resources
 ----------------

{pytorch_optimizer-2.9.1 → pytorch_optimizer-2.10.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytorch_optimizer"
-version = "2.9.1"
+version = "2.10.0"
 description = "optimizer & lr scheduler collections in PyTorch"
 license = "Apache-2.0"
 authors = ["kozistr <kozistr@gmail.com>"]
@@ -9,7 +9,7 @@ readme = "README.rst"
 homepage = "https://github.com/kozistr/pytorch_optimizer"
 repository = "https://github.com/kozistr/pytorch_optimizer"
 documentation = "https://pytorch-optimizers.readthedocs.io/en/latest"
-keywords = ["pytorch", "deep-learning", "optimizer", "lr scheduler", "A2Grad", "ASGD", "AccSGD", "AdaBelief", "AdaBound", "AdaDelta", "AdaFactor", "AdaMax", "AdaMod", "AdaNorm", "AdaPNM", "AdaSmooth", "Adai", "AdamP", "AdamS", "Adan", "AggMo", "AliG", "Apollo", "AvaGrad", "DAdaptAdaGrad", "DAdaptAdam", "DAdaptAdan", "DAdaptSGD", "DiffGrad", "Fromage", "Gravity", "LARS", "Lamb", "Lion", "MADGRAD", "MSVAG", "Nero", "NovoGrad", "PID", "PNM", "QHAdam", "QHM", "RAdam", "Ranger", "Ranger21", "SGDP", "SGDW", "SM3", "SRMM", "SWATS", "ScalableShampoo", "Shampoo", "Yogi", "SAM", "GSAM", "PCGrad", "RotoGrad"]
+keywords = ["pytorch", "deep-learning", "optimizer", "lr scheduler", "A2Grad", "ASGD", "AccSGD", "AdaBelief", "AdaBound", "AdaDelta", "AdaFactor", "AdaMax", "AdaMod", "AdaNorm", "AdaPNM", "AdaSmooth", "AdaHessian", "Adai", "AdamP", "AdamS", "Adan", "AggMo", "AliG", "Amos", "Apollo", "AvaGrad", "DAdaptAdaGrad", "DAdaptAdam", "DAdaptAdan", "DAdaptSGD", "DiffGrad", "Fromage", "Gravity", "LARS", "Lamb", "Lion", "MADGRAD", "MSVAG", "Nero", "NovoGrad", "PID", "PNM", "QHAdam", "QHM", "RAdam", "Ranger", "Ranger21", "SGDP", "SGDW", "SignSGD", "SM3", "SopihaH", "SRMM", "SWATS", "ScalableShampoo", "Shampoo", "Yogi", "SAM", "GSAM", "PCGrad", "RotoGrad"]
 classifiers = [
     "License :: OSI Approved :: Apache Software License",
     "Development Status :: 5 - Production/Stable",

{pytorch_optimizer-2.9.1 → pytorch_optimizer-2.10.0}/pytorch_optimizer/__init__.py

@@ -21,6 +21,7 @@ from pytorch_optimizer.optimizer.adabelief import AdaBelief
 from pytorch_optimizer.optimizer.adabound import AdaBound
 from pytorch_optimizer.optimizer.adadelta import AdaDelta
 from pytorch_optimizer.optimizer.adafactor import AdaFactor
+from pytorch_optimizer.optimizer.adahessian import AdaHessian
 from pytorch_optimizer.optimizer.adai import Adai
 from pytorch_optimizer.optimizer.adamax import AdaMax
 from pytorch_optimizer.optimizer.adamod import AdaMod
@@ -34,6 +35,7 @@ from pytorch_optimizer.optimizer.adasmooth import AdaSmooth
 from pytorch_optimizer.optimizer.agc import agc
 from pytorch_optimizer.optimizer.aggmo import AggMo
 from pytorch_optimizer.optimizer.alig import AliG
+from pytorch_optimizer.optimizer.amos import Amos
 from pytorch_optimizer.optimizer.apollo import Apollo
 from pytorch_optimizer.optimizer.avagrad import AvaGrad
 from pytorch_optimizer.optimizer.dadapt import DAdaptAdaGrad, DAdaptAdam, DAdaptAdan, DAdaptSGD
@@ -61,7 +63,7 @@ from pytorch_optimizer.optimizer.ranger import Ranger
 from pytorch_optimizer.optimizer.ranger21 import Ranger21
 from pytorch_optimizer.optimizer.rotograd import RotoGrad
 from pytorch_optimizer.optimizer.sam import SAM
-from pytorch_optimizer.optimizer.sgd import ASGD, SGDW, AccSGD
+from pytorch_optimizer.optimizer.sgd import ASGD, SGDW, AccSGD, SignSGD
 from pytorch_optimizer.optimizer.sgdp import SGDP
 from pytorch_optimizer.optimizer.shampoo import ScalableShampoo, Shampoo
 from pytorch_optimizer.optimizer.shampoo_utils import (
@@ -80,6 +82,7 @@ from pytorch_optimizer.optimizer.shampoo_utils import (
     power_iteration,
 )
 from pytorch_optimizer.optimizer.sm3 import SM3
+from pytorch_optimizer.optimizer.sophia import SophiaH
 from pytorch_optimizer.optimizer.srmm import SRMM
 from pytorch_optimizer.optimizer.swats import SWATS
 from pytorch_optimizer.optimizer.utils import (
@@ -145,6 +148,10 @@ OPTIMIZER_LIST: List[OPTIMIZER] = [
     AvaGrad,
     AdaShift,
     AdaDelta,
+    Amos,
+    AdaHessian,
+    SophiaH,
+    SignSGD,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
 

{pytorch_optimizer-2.9.1 → pytorch_optimizer-2.10.0}/pytorch_optimizer/base/optimizer.py

@@ -5,12 +5,80 @@ from typing import List, Optional, Tuple, Union
 import torch
 
 from pytorch_optimizer.base.exception import NegativeLRError, NegativeStepError
-from pytorch_optimizer.base.types import BETAS
+from pytorch_optimizer.base.types import BETAS, HUTCHINSON_G, PARAMETERS, STATE
 
 
 class BaseOptimizer(ABC):
     r"""Base optimizer class."""
 
+    @staticmethod
+    @torch.no_grad()
+    def set_hessian(param_groups: PARAMETERS, state: STATE, hessian: List[torch.Tensor]):
+        r"""Set hessian to state from external source. Generally useful when using functorch as a base.
+
+        Example usage:
+        ```
+        # Hutchinsons Estimator using HVP
+        noise = tree_map(lambda v: torch.randn_like(v), params)
+        loss_, hvp_est = jvp(grad(run_model_fn), (params,), (noise,))
+        hessian_diag_est  = tree_map(lambda a, b: a * b, hvp_est, noise)
+
+        optimizer.set_hessian(hessian_diag_est)
+        # OR
+        optimizer.step(hessian=hessian_diag_est)
+        ````
+        """
+        i: int = 0
+        for group in param_groups:
+            for p in group['params']:
+                if p.size() != hessian[i].size():
+                    raise ValueError(
+                        f'[-] the shape of parameter and hessian does not match. {p.size()} vs {hessian[i].size()}'
+                    )
+
+                state[p]['hessian'] = hessian[i]
+                i += 1
+
+    @staticmethod
+    @torch.no_grad()
+    def compute_hutchinson_hessian(
+        param_groups: PARAMETERS,
+        state: STATE,
+        num_samples: int = 1,
+        pre_zero: bool = True,
+        alpha: float = 1.0,
+        distribution: HUTCHINSON_G = 'gaussian',
+    ):
+        r"""Hutchinson's approximate hessian, added to the state under key `hessian`."""
+        if distribution not in ('gaussian', 'rademacher'):
+            raise NotImplementedError(f'[-] Hessian with distribution {distribution} is not implemented.')
+
+        params = []
+        for group in param_groups:
+            for p in group['params']:
+                if p.requires_grad and p.grad is not None and not p.grad.is_sparse:
+                    if 'hessian' not in state[p]:
+                        state[p]['hessian'] = torch.zeros_like(p)
+                    elif pre_zero:
+                        state[p]['hessian'].zero_()
+
+                    params.append(p)
+
+        if len(params) == 0:
+            return
+
+        grads = [p.grad for p in params]
+
+        for i in range(num_samples):
+            if distribution == 'rademacher':
+                zs = [torch.randint_like(p, 0, 1) * 2.0 - 1.0 for p in params]
+            else:
+                zs = [torch.randn_like(p) for p in params]
+
+            h_zs = torch.autograd.grad(grads, params, grad_outputs=zs, retain_graph=i < num_samples - 1)
+            for h_z, z, p in zip(h_zs, zs, params):
+                state[p]['hessian'].add_(h_z * z, alpha=(1 / num_samples) * alpha)
+
     @staticmethod
     def apply_weight_decay(
         p: torch.Tensor,

{pytorch_optimizer-2.9.1 → pytorch_optimizer-2.10.0}/pytorch_optimizer/base/types.py

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Dict, Iterable, Optional, Tuple, Type, Union
+from typing import Any, Callable, Dict, Iterable, Literal, Optional, Tuple, Type, Union
 
 import torch
 from torch.optim import Optimizer
@@ -12,3 +12,5 @@ PARAMETERS = Optional[Union[Iterable[Dict[str, Any]], Iterable[torch.Tensor]]]
 STATE = Dict[str, Any]
 OPTIMIZER = Type[Optimizer]
 SCHEDULER = Type[_LRScheduler]
+
+HUTCHINSON_G = Literal['gaussian', 'rademacher']

pytorch_optimizer-2.10.0/pytorch_optimizer/optimizer/adahessian.py

@@ -0,0 +1,139 @@
+from typing import List, Optional
+
+import torch
+from torch.optim.optimizer import Optimizer
+
+from pytorch_optimizer.base.exception import NoSparseGradientError
+from pytorch_optimizer.base.optimizer import BaseOptimizer
+from pytorch_optimizer.base.types import BETAS, CLOSURE, DEFAULTS, HUTCHINSON_G, LOSS, PARAMETERS
+
+
+class AdaHessian(Optimizer, BaseOptimizer):
+    r"""An Adaptive Second Order Optimizer for Machine Learning.
+
+        Requires `loss.backward(create_graph=True)` in order to calculate hessians.
+
+    :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
+    :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 hessian_power: float. exponent of the hessian trace.
+    :param update_period: int. number of steps after which to apply hessian approximation.
+    :param num_samples: int. times to sample `z` for the approximation of the hessian trace.
+    :param hessian_distribution: HUTCHINSON_G. type of distribution to initialize hessian.
+    :param adam_debias: bool. Only correct the denominator to avoid inflating step sizes early in training.
+    :param eps: float. term added to the denominator to improve numerical stability.
+    """
+
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 1e-1,
+        betas: BETAS = (0.9, 0.999),
+        weight_decay: float = 0.0,
+        weight_decouple: bool = True,
+        fixed_decay: bool = False,
+        hessian_power: float = 1.0,
+        update_period: int = 1,
+        num_samples: int = 1,
+        hessian_distribution: HUTCHINSON_G = 'rademacher',
+        adam_debias: bool = False,
+        eps: float = 1e-16,
+    ):
+        self.validate_learning_rate(lr)
+        self.validate_betas(betas)
+        self.validate_non_negative(weight_decay, 'weight_decay')
+        self.validate_non_negative(eps, 'eps')
+        self.validate_range(hessian_power, 'Hessian Power', 0, 1, range_type='(]')
+
+        self.update_period = update_period
+        self.num_samples = num_samples
+        self.distribution = hessian_distribution
+
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'betas': betas,
+            'weight_decay': weight_decay,
+            'weight_decouple': weight_decouple,
+            'fixed_decay': fixed_decay,
+            'hessian_power': hessian_power,
+            'adam_debias': adam_debias,
+            'eps': eps,
+        }
+        super().__init__(params, defaults)
+
+    def __str__(self) -> str:
+        return 'AdaHessian'
+
+    @torch.no_grad()
+    def reset(self):
+        for group in self.param_groups:
+            group['step'] = 0
+            for p in group['params']:
+                state = self.state[p]
+                state['exp_avg'] = torch.zeros_like(p)
+                state['exp_hessian_diag_sq'] = torch.zeros_like(p)
+
+    @torch.no_grad()
+    def step(self, closure: CLOSURE = None, hessian: Optional[List[torch.Tensor]] = None) -> LOSS:
+        loss: LOSS = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        step: int = self.param_groups[0].get('step', 1)
+
+        if hessian is not None:
+            self.set_hessian(self.param_groups, self.state, hessian)
+        elif step % self.update_period == 0:
+            self.compute_hutchinson_hessian(
+                self.param_groups, self.state, self.num_samples, distribution=self.distribution
+            )
+
+        for group in self.param_groups:
+            if 'step' in group:
+                group['step'] += 1
+            else:
+                group['step'] = 1
+
+            beta1, beta2 = group['betas']
+
+            bias_correction1: float = 1.0 - beta1 ** group['step']
+            bias_correction2: float = 1.0 - beta2 ** group['step']
+
+            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 'exp_avg' not in state:
+                    state['exp_avg'] = torch.zeros_like(p)
+                    state['exp_hessian_diag_sq'] = torch.zeros_like(p)
+
+                self.apply_weight_decay(
+                    p=p,
+                    grad=grad,
+                    lr=group['lr'],
+                    weight_decay=group['weight_decay'],
+                    weight_decouple=group['weight_decouple'],
+                    fixed_decay=group['fixed_decay'],
+                )
+
+                exp_avg, exp_hessian_diag_sq = state['exp_avg'], state['exp_hessian_diag_sq']
+                exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)
+
+                if 'hessian' in state and (group['step'] % self.update_period == 0 or hessian is not None):
+                    exp_hessian_diag_sq.mul_(beta2).addcmul_(state['hessian'], state['hessian'], value=1.0 - beta2)
+
+                de_nom = (exp_hessian_diag_sq / bias_correction2).pow_(group['hessian_power'] / 2).add_(group['eps'])
+
+                step_size: float = self.apply_adam_debias(group['adam_debias'], group['lr'], bias_correction1)
+                p.addcdiv_(exp_avg, de_nom, value=-step_size)
+
+        return loss

pytorch_optimizer-2.10.0/pytorch_optimizer/optimizer/amos.py

@@ -0,0 +1,142 @@
+import math
+
+import torch
+from torch.optim.optimizer import Optimizer
+
+from pytorch_optimizer.base.exception import NoSparseGradientError
+from pytorch_optimizer.base.optimizer import BaseOptimizer
+from pytorch_optimizer.base.types import CLOSURE, DEFAULTS, LOSS, PARAMETERS
+
+
+class Amos(Optimizer, BaseOptimizer):
+    r"""An Adam-style Optimizer with Adaptive Weight Decay towards Model-Oriented Scale.
+
+    :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
+    :param lr: float. learning rate.
+    :param beta: float. A float slightly < 1. We recommend setting `1 - beta` to the same order of magnitude
+        as the learning rate. similarity with beta2 in Adam.
+    :param momentum: float. Exponential decay rate for optional moving average of updates.
+    :param extra_l2: float. Additional L2 regularization.
+    :param c_coef: float. Coefficient for decay_factor_c.
+    :param d_coef: float. Coefficient for decay_factor_d.
+    :param eps: float. term added to the denominator to improve numerical stability.
+    """
+
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 1e-3,
+        beta: float = 0.999,
+        momentum: float = 0.0,
+        extra_l2: float = 0.0,
+        c_coef: float = 0.25,
+        d_coef: float = 0.25,
+        eps: float = 1e-18,
+    ):
+        self.validate_learning_rate(lr)
+        self.validate_range(momentum, 'momentum', 0.0, 1.0, range_type='[)')
+        self.validate_range(beta, 'beta', 0.0, 1.0, range_type='[)')
+        self.validate_non_negative(extra_l2, 'extra_l2')
+        self.validate_non_negative(eps, 'eps')
+
+        self.c_coef = c_coef
+        self.d_coef = d_coef
+
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'beta': beta,
+            'momentum': momentum,
+            'extra_l2': extra_l2,
+            'eps': eps,
+        }
+
+        super().__init__(params, defaults)
+
+    def __str__(self) -> str:
+        return 'Amos'
+
+    @torch.no_grad()
+    def reset(self):
+        for group in self.param_groups:
+            group['step'] = 0
+            for p in group['params']:
+                state = self.state[p]
+
+                state['exp_avg_sq'] = torch.zeros((1,), dtype=p.dtype, device=p.device)
+                state['decay'] = torch.zeros((1,), dtype=p.dtype, device=p.device)
+                if group['momentum'] > 0.0:
+                    state['exp_avg'] = torch.zeros_like(p)
+
+    @staticmethod
+    def get_scale(p: torch.Tensor) -> float:
+        r"""Get expected scale for model weights."""
+        if len(p.shape) == 1:  # expected 'bias'
+            return 0.5
+        if len(p.shape) == 2:  # expected Embedding, Linear, ...
+            return math.sqrt(2 / p.size(1))
+        return math.sqrt(1 / p.size(1))
+
+    @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:
+            if 'step' in group:
+                group['step'] += 1
+            else:
+                group['step'] = 1
+
+            momentum, beta = group['momentum'], group['beta']
+
+            lr_sq: float = math.sqrt(group['lr'])
+            bias_correction: float = 1.0 - beta ** group['step']
+
+            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 len(state) == 0:
+                    state['exp_avg_sq'] = torch.zeros((1,), dtype=p.dtype, device=p.device)
+                    state['decay'] = torch.zeros((1,), dtype=p.dtype, device=p.device)
+                    if group['momentum'] > 0.0:
+                        state['exp_avg'] = torch.zeros_like(p)
+
+                g2 = grad.pow(2).mean()
+                init_lr: float = group['lr'] * self.get_scale(p)
+
+                exp_avg_sq = state['exp_avg_sq']
+                exp_avg_sq.mul_(beta).add_(g2, alpha=1.0 - beta)
+
+                r_v_hat = bias_correction / (exp_avg_sq + group['eps'])
+
+                b = state['decay']
+                decay_factor_c = torch.rsqrt(1.0 + self.c_coef * lr_sq * b)
+                decay_factor_d = torch.reciprocal(1.0 + self.d_coef * math.sqrt(init_lr) * b)
+
+                gamma = decay_factor_c * (group['lr'] ** 2) * r_v_hat * g2
+
+                update = p.clone()
+                update.mul_((gamma - group['extra_l2']) / 2.0)
+                update.add_(r_v_hat.sqrt() * grad, alpha=init_lr)
+                update.mul_(decay_factor_d)
+
+                b.mul_(1.0 + gamma).add_(gamma)
+
+                if momentum > 0.0:
+                    exp_avg = state['exp_avg']
+                    exp_avg.mul_(momentum).add_(update, alpha=1.0 - momentum)
+
+                    update.copy_(exp_avg)
+
+                p.add_(-update)
+
+        return loss

{pytorch_optimizer-2.9.1 → pytorch_optimizer-2.10.0}/pytorch_optimizer/optimizer/dadapt.py

@@ -266,7 +266,7 @@ class DAdaptAdam(Optimizer, BaseOptimizer):
         weight_decouple: bool = False,
         fixed_decay: bool = False,
         bias_correction: bool = False,
-        eps: float = 0.0,
+        eps: float = 1e-8,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-2.9.1 → pytorch_optimizer-2.10.0}/pytorch_optimizer/optimizer/sgd.py

@@ -311,3 +311,75 @@ class ASGD(Optimizer, BaseOptimizer):
                 p.add_(grad, alpha=-new_lr)
 
         return loss
+
+
+class SignSGD(Optimizer, BaseOptimizer):
+    r"""Compressed Optimisation for Non-Convex Problems.
+
+    :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
+    :param lr: float. learning rate.
+    :param momentum: float. momentum factor (0.0 = SignSGD, >0 = Signum).
+    :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-3,
+        momentum: float = 0.9,
+        weight_decay: float = 0.0,
+        weight_decouple: bool = True,
+    ):
+        self.validate_learning_rate(lr)
+        self.validate_range(momentum, 'beta', 0.0, 1.0)
+        self.validate_non_negative(weight_decay, 'weight_decay')
+
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'momentum': momentum,
+            'weight_decay': weight_decay,
+            'weight_decouple': weight_decouple,
+        }
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def reset(self):
+        for group in self.param_groups:
+            group['step'] = 0
+            for p in group['params']:
+                state = self.state[p]
+
+                if group['momentum'] > 0.0:
+                    state['momentum_buffer'] = torch.zeros_like(p)
+
+    @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:
+            momentum = group['momentum']
+            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 momentum > 0.0:
+                    if len(state) == 0:
+                        state['momentum_buffer'] = torch.zeros_like(p)
+
+                    buf = state['momentum_buffer']
+                    buf.mul_(momentum).add_(grad, alpha=1.0 - momentum)
+                else:
+                    buf = grad
+
+                p.add_(torch.sign(buf), alpha=-group['lr'])
+
+        return loss

pytorch_optimizer-2.10.0/pytorch_optimizer/optimizer/sophia.py

@@ -0,0 +1,133 @@
+from typing import List, Optional
+
+import torch
+from torch.optim.optimizer import Optimizer
+
+from pytorch_optimizer.base.exception import NoSparseGradientError
+from pytorch_optimizer.base.optimizer import BaseOptimizer
+from pytorch_optimizer.base.types import BETAS, CLOSURE, DEFAULTS, HUTCHINSON_G, LOSS, PARAMETERS
+
+
+class SophiaH(Optimizer, BaseOptimizer):
+    r"""Second-order Clipped Stochastic Optimization.
+
+        Requires `loss.backward(create_graph=True)` in order to calculate hessians.
+
+    :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
+    :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 p: float. clip effective (applied) gradient (p).
+    :param update_period: int. number of steps after which to apply hessian approximation.
+    :param num_samples: int. times to sample `z` for the approximation of the hessian trace.
+    :param hessian_distribution: HUTCHINSON_G. type of distribution to initialize hessian.
+    :param eps: float. term added to the denominator to improve numerical stability.
+    """
+
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 6e-2,
+        betas: BETAS = (0.96, 0.99),
+        weight_decay: float = 0.0,
+        weight_decouple: bool = True,
+        fixed_decay: bool = False,
+        p: float = 1e-2,
+        update_period: int = 10,
+        num_samples: int = 1,
+        hessian_distribution: HUTCHINSON_G = 'gaussian',
+        eps: float = 1e-12,
+    ):
+        self.validate_learning_rate(lr)
+        self.validate_betas(betas)
+        self.validate_non_negative(weight_decay, 'weight_decay')
+        self.validate_non_negative(p, 'p (gradient clip)')
+        self.validate_step(update_period, 'update_period')
+        self.validate_positive(num_samples, 'num_samples')
+        self.validate_options(hessian_distribution, 'hessian_distribution', ['gaussian', 'rademacher'])
+        self.validate_non_negative(eps, 'eps')
+
+        self.update_period = update_period
+        self.num_samples = num_samples
+        self.distribution = hessian_distribution
+
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'betas': betas,
+            'weight_decay': weight_decay,
+            'weight_decouple': weight_decouple,
+            'fixed_decay': fixed_decay,
+            'p': p,
+            'eps': eps,
+        }
+        super().__init__(params, defaults)
+
+    def __str__(self) -> str:
+        return 'SophiaH'
+
+    @torch.no_grad()
+    def reset(self):
+        for group in self.param_groups:
+            group['step'] = 0
+            for p in group['params']:
+                state = self.state[p]
+                state['momentum'] = torch.zeros_like(p)
+                state['hessian_moment'] = torch.zeros_like(p)
+
+    @torch.no_grad()
+    def step(self, closure: CLOSURE = None, hessian: Optional[List[torch.Tensor]] = None) -> LOSS:
+        loss: LOSS = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        step: int = self.param_groups[0].get('step', 1)
+
+        if hessian is not None:
+            self.set_hessian(self.param_groups, self.state, hessian)
+        elif step % self.update_period == 0:
+            self.compute_hutchinson_hessian(
+                self.param_groups, self.state, self.num_samples, distribution=self.distribution
+            )
+
+        for group in self.param_groups:
+            if 'step' in group:
+                group['step'] += 1
+            else:
+                group['step'] = 1
+
+            beta1, beta2 = group['betas']
+            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 len(state) == 0:
+                    state['momentum'] = torch.zeros_like(p)
+                    state['hessian_moment'] = torch.zeros_like(p)
+
+                self.apply_weight_decay(
+                    p=p,
+                    grad=grad,
+                    lr=group['lr'],
+                    weight_decay=group['weight_decay'],
+                    weight_decouple=group['weight_decouple'],
+                    fixed_decay=group['fixed_decay'],
+                )
+
+                momentum, hessian_moment = state['momentum'], state['hessian_moment']
+                momentum.mul_(beta1).add_(grad, alpha=1.0 - beta1)
+
+                if 'hessian' in state and (group['step'] % self.update_period == 0 or hessian is not None):
+                    hessian_moment.mul_(beta2).add_(state['hessian'], alpha=1.0 - beta2)
+
+                update = (momentum / torch.clip(hessian_moment, min=group['eps'])).clamp_(-group['p'], group['p'])
+                p.add_(update, alpha=-group['lr'])
+
+        return loss